Index: appengine/swarming/server/bot_code_test.py |
diff --git a/appengine/swarming/server/bot_code_test.py b/appengine/swarming/server/bot_code_test.py |
index 46218e519aea147fe0f27195af850d718d31d1c7..854500eb0da9d13c372cc835e15ecb254f6089dc 100755 |
--- a/appengine/swarming/server/bot_code_test.py |
+++ b/appengine/swarming/server/bot_code_test.py |
@@ -10,6 +10,7 @@ import re |
import subprocess |
import sys |
import tempfile |
+import time |
import unittest |
import zipfile |
@@ -40,6 +41,7 @@ class BotManagementTest(test_case.TestCase): |
auth, 'get_current_identity', |
lambda: auth.Identity(auth.IDENTITY_USER, 'joe@localhost')) |
+ |
def test_get_bootstrap(self): |
def get_self_config_mock(path, revision=None, store_last_good=False): |
self.assertEqual('scripts/bootstrap.py', path) |
@@ -94,7 +96,37 @@ class BotManagementTest(test_case.TestCase): |
self.assertEqual(expected, additionals) |
def test_get_swarming_bot_zip(self): |
+ local_mc = {'store': {}, 'reads': 0, 'writes': 0} |
+ |
+ @bot_code.ndb.tasklet |
Vadim Sh.
2017/06/28 17:52:39
please import ndb directly in this module (somewhe
aludwin
2017/06/28 19:19:49
Done.
|
+ def mock_memcache_get(version, desc, part=None): |
+ value = local_mc['store'].get(bot_code.bot_key(version, desc, part)) |
+ if value is not None: |
+ local_mc['reads'] += 1 |
+ raise bot_code.ndb.Return(value) |
+ |
+ @bot_code.ndb.tasklet |
+ def mock_memcache_set(value, version, desc, part=None): |
+ local_mc['writes'] += 1 |
+ key = bot_code.bot_key(version, desc, part) |
+ local_mc['store'][key] = value |
+ return bot_code.ndb.Return(None) |
+ |
+ self.mock(bot_code, 'bot_memcache_set', mock_memcache_set) |
+ self.mock(bot_code, 'bot_memcache_get', mock_memcache_get) |
+ self.mock(bot_code, 'MAX_MEMCACHED_SIZE_BYTES', 100000) |
+ |
+ self.assertEqual(0, local_mc['writes']) |
zipped_code = bot_code.get_swarming_bot_zip('http://localhost') |
+ self.assertEqual(0, local_mc['reads']) |
+ self.assertNotEqual(0, local_mc['writes']) |
+ |
+ # Make sure that we read from memcached if we get it again |
+ zipped_code_copy = bot_code.get_swarming_bot_zip('http://localhost') |
+ self.assertEqual(local_mc['writes'], local_mc['reads']) |
+ # 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: |
@@ -120,6 +152,7 @@ class BotManagementTest(test_case.TestCase): |
finally: |
file_path.rmtree(temp_dir) |
+ |
def test_bootstrap_token(self): |
tok = bot_code.generate_bootstrap_token() |
self.assertEqual( |