Replace custom blob gRPC API with ByteStream

appengine/swarming/server/bot_code_test.py

 #!/usr/bin/env python
 # Copyright 2014 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 import StringIO
 import logging
 import os
 import re
 import subprocess
@@ skipping 22 matching lines @@
 
 class BotManagementTest(test_case.TestCase):
   def setUp(self):
     super(BotManagementTest, self).setUp()
     self.testbed.init_user_stub()
 
     self.mock(
         auth, 'get_current_identity',
         lambda: auth.Identity(auth.IDENTITY_USER, 'joe@localhost'))
 
+    self._memcached={}
+    self._read_from_memcached = False
+
+    def mock_set_cached_bot_entry(content, version, desc, part=None):
+      key = bot_code.cached_bot_key(version, desc, part)
+      self._memcached[key] = content
+
+    def mock_get_cached_bot_entry(version, desc, part=None):
+      key = bot_code.cached_bot_key(version, desc, part)
+      if key not in self._memcached:
+        raise bot_code.MemcacheMissingException
+      self._read_from_memcached = True
+      return self._memcached[key]
+
+    self.mock(bot_code, 'set_cached_bot_entry', mock_set_cached_bot_entry)
+    self.mock(bot_code, 'get_cached_bot_entry', mock_get_cached_bot_entry)
+    self.mock(bot_code, 'MAX_MEMCACHED_SIZE_BYTES', 100000)
+
   def test_get_bootstrap(self):
     def get_self_config_mock(path, revision=None, store_last_good=False):
       self.assertEqual('scripts/bootstrap.py', path)
       self.assertEqual(None, revision)
       self.assertEqual(True, store_last_good)
       return None, 'foo bar'
     self.mock(config, 'get_self_config', get_self_config_mock)
     f = bot_code.get_bootstrap('localhost', 'token', None)
     expected = (
       '#!/usr/bin/env python\n'
@@ skipping 35 matching lines @@
   def test_get_bot_version(self):
     actual, additionals = bot_code.get_bot_version('http://localhost')
     self.assertTrue(re.match(r'^[0-9a-f]{64}$', actual), actual)
     expected = {
       'config/bot_config.py': bot_code.get_bot_config().content,
     }
     self.assertEqual(expected, additionals)
 
   def test_get_swarming_bot_zip(self):
     zipped_code = bot_code.get_swarming_bot_zip('http://localhost')
+    self.assertFalse(self._read_from_memcached)
+    # Make sure that we read from memcached if we get it again
+    zipped_code_copy = bot_code.get_swarming_bot_zip('http://localhost')
+    self.assertTrue(self._read_from_memcached)
+    # Why not assertEqual? Don't want to dump ~1MB of data if this fails.
+    self.assertTrue(zipped_code == zipped_code_copy)
     # Ensure the zip is valid and all the expected files are present.
     with zipfile.ZipFile(StringIO.StringIO(zipped_code), 'r') as zip_file:
       for i in bot_archive.FILES:
         with zip_file.open(i) as f:
           content = f.read()
           if os.path.basename(i) != '__init__.py':
             self.assertTrue(content, i)
 
     temp_dir = tempfile.mkdtemp(prefix='swarming')
     try:
       # Try running the bot and ensure it can import the required files. (It
       # would crash if it failed to import them).
       bot_path = os.path.join(temp_dir, 'swarming_bot.zip')
       with open(bot_path, 'wb') as f:
         f.write(zipped_code)
       proc = subprocess.Popen(
           [sys.executable, bot_path, 'start_bot', '-h'],
           cwd=temp_dir,
           stdout=subprocess.PIPE,
           stderr=subprocess.STDOUT)
       out = proc.communicate()[0]
       self.assertEqual(0, proc.returncode, out)
     finally:
       file_path.rmtree(temp_dir)
 
+
   def test_bootstrap_token(self):
     tok = bot_code.generate_bootstrap_token()
     self.assertEqual(
         {'for': 'user:joe@localhost'}, bot_code.validate_bootstrap_token(tok))
 
 
 if __name__ == '__main__':
   fix_encoding.fix_encoding()
   logging.basicConfig(
       level=logging.DEBUG if '-v' in sys.argv else logging.ERROR)
   if '-v' in sys.argv:
     unittest.TestCase.maxDiff = None
   unittest.main()