OLD | NEW |
---|---|
1 #!/usr/bin/env python | 1 #!/usr/bin/env python |
2 # Copyright 2014 The LUCI Authors. All rights reserved. | 2 # Copyright 2014 The LUCI Authors. All rights reserved. |
3 # Use of this source code is governed under the Apache License, Version 2.0 | 3 # Use of this source code is governed under the Apache License, Version 2.0 |
4 # that can be found in the LICENSE file. | 4 # that can be found in the LICENSE file. |
5 | 5 |
6 import StringIO | 6 import StringIO |
7 import logging | 7 import logging |
8 import os | 8 import os |
9 import re | 9 import re |
10 import subprocess | 10 import subprocess |
11 import sys | 11 import sys |
12 import tempfile | 12 import tempfile |
13 import time | |
13 import unittest | 14 import unittest |
14 import zipfile | 15 import zipfile |
15 | 16 |
16 import test_env | 17 import test_env |
17 test_env.setup_test_env() | 18 test_env.setup_test_env() |
18 | 19 |
19 from components import auth | 20 from components import auth |
20 from test_support import test_case | 21 from test_support import test_case |
21 | 22 |
22 from server import bot_archive | 23 from server import bot_archive |
(...skipping 10 matching lines...) Expand all Loading... | |
33 | 34 |
34 class BotManagementTest(test_case.TestCase): | 35 class BotManagementTest(test_case.TestCase): |
35 def setUp(self): | 36 def setUp(self): |
36 super(BotManagementTest, self).setUp() | 37 super(BotManagementTest, self).setUp() |
37 self.testbed.init_user_stub() | 38 self.testbed.init_user_stub() |
38 | 39 |
39 self.mock( | 40 self.mock( |
40 auth, 'get_current_identity', | 41 auth, 'get_current_identity', |
41 lambda: auth.Identity(auth.IDENTITY_USER, 'joe@localhost')) | 42 lambda: auth.Identity(auth.IDENTITY_USER, 'joe@localhost')) |
42 | 43 |
44 | |
43 def test_get_bootstrap(self): | 45 def test_get_bootstrap(self): |
44 def get_self_config_mock(path, revision=None, store_last_good=False): | 46 def get_self_config_mock(path, revision=None, store_last_good=False): |
45 self.assertEqual('scripts/bootstrap.py', path) | 47 self.assertEqual('scripts/bootstrap.py', path) |
46 self.assertEqual(None, revision) | 48 self.assertEqual(None, revision) |
47 self.assertEqual(True, store_last_good) | 49 self.assertEqual(True, store_last_good) |
48 return None, 'foo bar' | 50 return None, 'foo bar' |
49 self.mock(config, 'get_self_config', get_self_config_mock) | 51 self.mock(config, 'get_self_config', get_self_config_mock) |
50 f = bot_code.get_bootstrap('localhost', 'token', None) | 52 f = bot_code.get_bootstrap('localhost', 'token', None) |
51 expected = ( | 53 expected = ( |
52 '#!/usr/bin/env python\n' | 54 '#!/usr/bin/env python\n' |
(...skipping 34 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... | |
87 | 89 |
88 def test_get_bot_version(self): | 90 def test_get_bot_version(self): |
89 actual, additionals = bot_code.get_bot_version('http://localhost') | 91 actual, additionals = bot_code.get_bot_version('http://localhost') |
90 self.assertTrue(re.match(r'^[0-9a-f]{64}$', actual), actual) | 92 self.assertTrue(re.match(r'^[0-9a-f]{64}$', actual), actual) |
91 expected = { | 93 expected = { |
92 'config/bot_config.py': bot_code.get_bot_config().content, | 94 'config/bot_config.py': bot_code.get_bot_config().content, |
93 } | 95 } |
94 self.assertEqual(expected, additionals) | 96 self.assertEqual(expected, additionals) |
95 | 97 |
96 def test_get_swarming_bot_zip(self): | 98 def test_get_swarming_bot_zip(self): |
99 local_mc = {'store': {}, 'reads': 0, 'writes': 0} | |
100 | |
101 def mock_memcache_get(version, desc): | |
102 result = local_mc['store'].get(bot_code.bot_key(version, desc)) | |
103 if result is not None: | |
104 local_mc['reads'] += 1 | |
105 return result | |
106 | |
107 def mock_memcache_async_get(version, desc, part=None): | |
Vadim Sh.
2017/06/27 19:12:49
this will be more correct mock:
@ndb.tasklet
def
aludwin
2017/06/27 20:10:17
How would I do the counting in this case? I'd like
Vadim Sh.
2017/06/27 21:14:29
yeah, this will not be possible.
But why counting
aludwin
2017/06/28 13:23:53
Just thought it would be nice to ensure that every
| |
108 class Future: | |
109 def __init__(self, value): | |
110 self._value = value | |
111 def get_result(self): | |
112 if self._value is not None: | |
113 local_mc['reads'] += 1 | |
114 return self._value | |
115 value = local_mc['store'].get(bot_code.bot_key(version, desc, part)) | |
116 return Future(value) | |
117 | |
118 def mock_memcache_async_set(value, version, desc, part=None): | |
119 class Future: | |
120 def __init__(self, key, value): | |
121 self._key = key | |
122 self._value = value | |
123 @property | |
124 def state(self): | |
125 local_mc['writes'] += 1 | |
126 local_mc['store'][key] = value | |
127 return bot_code.ndb.Future.FINISHING | |
128 key = bot_code.bot_key(version, desc, part) | |
129 return Future(key, value) | |
130 | |
131 self.mock(bot_code, 'bot_memcache_async_set', mock_memcache_async_set) | |
132 self.mock(bot_code, 'bot_memcache_async_get', mock_memcache_async_get) | |
133 self.mock(bot_code, 'bot_memcache_get', mock_memcache_get) | |
134 self.mock(bot_code, 'MAX_MEMCACHED_SIZE_BYTES', 100000) | |
135 | |
136 self.assertEqual(0, local_mc['writes']) | |
97 zipped_code = bot_code.get_swarming_bot_zip('http://localhost') | 137 zipped_code = bot_code.get_swarming_bot_zip('http://localhost') |
138 self.assertEqual(0, local_mc['reads']) | |
139 self.assertNotEqual(0, local_mc['writes']) | |
140 | |
141 # Make sure that we read from memcached if we get it again | |
142 zipped_code_copy = bot_code.get_swarming_bot_zip('http://localhost') | |
143 self.assertEqual(local_mc['writes'], local_mc['reads']) | |
144 # Why not assertEqual? Don't want to dump ~1MB of data if this fails. | |
145 self.assertTrue(zipped_code == zipped_code_copy) | |
146 | |
98 # Ensure the zip is valid and all the expected files are present. | 147 # Ensure the zip is valid and all the expected files are present. |
99 with zipfile.ZipFile(StringIO.StringIO(zipped_code), 'r') as zip_file: | 148 with zipfile.ZipFile(StringIO.StringIO(zipped_code), 'r') as zip_file: |
100 for i in bot_archive.FILES: | 149 for i in bot_archive.FILES: |
101 with zip_file.open(i) as f: | 150 with zip_file.open(i) as f: |
102 content = f.read() | 151 content = f.read() |
103 if os.path.basename(i) != '__init__.py': | 152 if os.path.basename(i) != '__init__.py': |
104 self.assertTrue(content, i) | 153 self.assertTrue(content, i) |
105 | 154 |
106 temp_dir = tempfile.mkdtemp(prefix='swarming') | 155 temp_dir = tempfile.mkdtemp(prefix='swarming') |
107 try: | 156 try: |
108 # Try running the bot and ensure it can import the required files. (It | 157 # Try running the bot and ensure it can import the required files. (It |
109 # would crash if it failed to import them). | 158 # would crash if it failed to import them). |
110 bot_path = os.path.join(temp_dir, 'swarming_bot.zip') | 159 bot_path = os.path.join(temp_dir, 'swarming_bot.zip') |
111 with open(bot_path, 'wb') as f: | 160 with open(bot_path, 'wb') as f: |
112 f.write(zipped_code) | 161 f.write(zipped_code) |
113 proc = subprocess.Popen( | 162 proc = subprocess.Popen( |
114 [sys.executable, bot_path, 'start_bot', '-h'], | 163 [sys.executable, bot_path, 'start_bot', '-h'], |
115 cwd=temp_dir, | 164 cwd=temp_dir, |
116 stdout=subprocess.PIPE, | 165 stdout=subprocess.PIPE, |
117 stderr=subprocess.STDOUT) | 166 stderr=subprocess.STDOUT) |
118 out = proc.communicate()[0] | 167 out = proc.communicate()[0] |
119 self.assertEqual(0, proc.returncode, out) | 168 self.assertEqual(0, proc.returncode, out) |
120 finally: | 169 finally: |
121 file_path.rmtree(temp_dir) | 170 file_path.rmtree(temp_dir) |
122 | 171 |
172 | |
123 def test_bootstrap_token(self): | 173 def test_bootstrap_token(self): |
124 tok = bot_code.generate_bootstrap_token() | 174 tok = bot_code.generate_bootstrap_token() |
125 self.assertEqual( | 175 self.assertEqual( |
126 {'for': 'user:joe@localhost'}, bot_code.validate_bootstrap_token(tok)) | 176 {'for': 'user:joe@localhost'}, bot_code.validate_bootstrap_token(tok)) |
127 | 177 |
128 | 178 |
129 if __name__ == '__main__': | 179 if __name__ == '__main__': |
130 fix_encoding.fix_encoding() | 180 fix_encoding.fix_encoding() |
131 logging.basicConfig( | 181 logging.basicConfig( |
132 level=logging.DEBUG if '-v' in sys.argv else logging.ERROR) | 182 level=logging.DEBUG if '-v' in sys.argv else logging.ERROR) |
133 if '-v' in sys.argv: | 183 if '-v' in sys.argv: |
134 unittest.TestCase.maxDiff = None | 184 unittest.TestCase.maxDiff = None |
135 unittest.main() | 185 unittest.main() |
OLD | NEW |