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, part=None): |
| 102 class Future: |
| 103 def __init__(self, value): |
| 104 self._value = value |
| 105 def get_result(self): |
| 106 if self._value is not None: |
| 107 local_mc['reads'] += 1 |
| 108 return self._value |
| 109 value = local_mc['store'].get(bot_code.bot_key(version, desc, part)) |
| 110 return Future(value) |
| 111 |
| 112 def mock_memcache_set(value, version, desc, part=None): |
| 113 class Future: |
| 114 def __init__(self, key, value): |
| 115 self._key = key |
| 116 self._value = value |
| 117 def check_success(self): |
| 118 local_mc['writes'] += 1 |
| 119 local_mc['store'][key] = value |
| 120 key = bot_code.bot_key(version, desc, part) |
| 121 return Future(key, value) |
| 122 |
| 123 self.mock(bot_code, 'bot_memcache_set', mock_memcache_set) |
| 124 self.mock(bot_code, 'bot_memcache_get', mock_memcache_get) |
| 125 self.mock(bot_code, 'MAX_MEMCACHED_SIZE_BYTES', 100000) |
| 126 |
| 127 self.assertEqual(0, local_mc['writes']) |
97 zipped_code = bot_code.get_swarming_bot_zip('http://localhost') | 128 zipped_code = bot_code.get_swarming_bot_zip('http://localhost') |
| 129 self.assertEqual(0, local_mc['reads']) |
| 130 self.assertNotEqual(0, local_mc['writes']) |
| 131 |
| 132 # Make sure that we read from memcached if we get it again |
| 133 zipped_code_copy = bot_code.get_swarming_bot_zip('http://localhost') |
| 134 self.assertEqual(local_mc['writes'], local_mc['reads']) |
| 135 # Why not assertEqual? Don't want to dump ~1MB of data if this fails. |
| 136 self.assertTrue(zipped_code == zipped_code_copy) |
| 137 |
98 # Ensure the zip is valid and all the expected files are present. | 138 # Ensure the zip is valid and all the expected files are present. |
99 with zipfile.ZipFile(StringIO.StringIO(zipped_code), 'r') as zip_file: | 139 with zipfile.ZipFile(StringIO.StringIO(zipped_code), 'r') as zip_file: |
100 for i in bot_archive.FILES: | 140 for i in bot_archive.FILES: |
101 with zip_file.open(i) as f: | 141 with zip_file.open(i) as f: |
102 content = f.read() | 142 content = f.read() |
103 if os.path.basename(i) != '__init__.py': | 143 if os.path.basename(i) != '__init__.py': |
104 self.assertTrue(content, i) | 144 self.assertTrue(content, i) |
105 | 145 |
106 temp_dir = tempfile.mkdtemp(prefix='swarming') | 146 temp_dir = tempfile.mkdtemp(prefix='swarming') |
107 try: | 147 try: |
108 # Try running the bot and ensure it can import the required files. (It | 148 # Try running the bot and ensure it can import the required files. (It |
109 # would crash if it failed to import them). | 149 # would crash if it failed to import them). |
110 bot_path = os.path.join(temp_dir, 'swarming_bot.zip') | 150 bot_path = os.path.join(temp_dir, 'swarming_bot.zip') |
111 with open(bot_path, 'wb') as f: | 151 with open(bot_path, 'wb') as f: |
112 f.write(zipped_code) | 152 f.write(zipped_code) |
113 proc = subprocess.Popen( | 153 proc = subprocess.Popen( |
114 [sys.executable, bot_path, 'start_bot', '-h'], | 154 [sys.executable, bot_path, 'start_bot', '-h'], |
115 cwd=temp_dir, | 155 cwd=temp_dir, |
116 stdout=subprocess.PIPE, | 156 stdout=subprocess.PIPE, |
117 stderr=subprocess.STDOUT) | 157 stderr=subprocess.STDOUT) |
118 out = proc.communicate()[0] | 158 out = proc.communicate()[0] |
119 self.assertEqual(0, proc.returncode, out) | 159 self.assertEqual(0, proc.returncode, out) |
120 finally: | 160 finally: |
121 file_path.rmtree(temp_dir) | 161 file_path.rmtree(temp_dir) |
122 | 162 |
| 163 |
123 def test_bootstrap_token(self): | 164 def test_bootstrap_token(self): |
124 tok = bot_code.generate_bootstrap_token() | 165 tok = bot_code.generate_bootstrap_token() |
125 self.assertEqual( | 166 self.assertEqual( |
126 {'for': 'user:joe@localhost'}, bot_code.validate_bootstrap_token(tok)) | 167 {'for': 'user:joe@localhost'}, bot_code.validate_bootstrap_token(tok)) |
127 | 168 |
128 | 169 |
129 if __name__ == '__main__': | 170 if __name__ == '__main__': |
130 fix_encoding.fix_encoding() | 171 fix_encoding.fix_encoding() |
131 logging.basicConfig( | 172 logging.basicConfig( |
132 level=logging.DEBUG if '-v' in sys.argv else logging.ERROR) | 173 level=logging.DEBUG if '-v' in sys.argv else logging.ERROR) |
133 if '-v' in sys.argv: | 174 if '-v' in sys.argv: |
134 unittest.TestCase.maxDiff = None | 175 unittest.TestCase.maxDiff = None |
135 unittest.main() | 176 unittest.main() |
OLD | NEW |