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 |
| 20 from google.appengine.ext import ndb |
19 from components import auth | 21 from components import auth |
20 from test_support import test_case | 22 from test_support import test_case |
21 | 23 |
22 from server import bot_archive | 24 from server import bot_archive |
23 from server import bot_code | 25 from server import bot_code |
24 from components import config | 26 from components import config |
25 | 27 |
26 CLIENT_DIR = os.path.join( | 28 CLIENT_DIR = os.path.join( |
27 os.path.dirname(os.path.dirname(test_env.APP_DIR)), 'client') | 29 os.path.dirname(os.path.dirname(test_env.APP_DIR)), 'client') |
28 sys.path.insert(0, CLIENT_DIR) | 30 sys.path.insert(0, CLIENT_DIR) |
29 from third_party.depot_tools import fix_encoding | 31 from third_party.depot_tools import fix_encoding |
30 from utils import file_path | 32 from utils import file_path |
31 sys.path.pop(0) | 33 sys.path.pop(0) |
32 | 34 |
33 | 35 |
34 class BotManagementTest(test_case.TestCase): | 36 class BotManagementTest(test_case.TestCase): |
35 def setUp(self): | 37 def setUp(self): |
36 super(BotManagementTest, self).setUp() | 38 super(BotManagementTest, self).setUp() |
37 self.testbed.init_user_stub() | 39 self.testbed.init_user_stub() |
38 | 40 |
39 self.mock( | 41 self.mock( |
40 auth, 'get_current_identity', | 42 auth, 'get_current_identity', |
41 lambda: auth.Identity(auth.IDENTITY_USER, 'joe@localhost')) | 43 lambda: auth.Identity(auth.IDENTITY_USER, 'joe@localhost')) |
42 | 44 |
| 45 |
43 def test_get_bootstrap(self): | 46 def test_get_bootstrap(self): |
44 def get_self_config_mock(path, revision=None, store_last_good=False): | 47 def get_self_config_mock(path, revision=None, store_last_good=False): |
45 self.assertEqual('scripts/bootstrap.py', path) | 48 self.assertEqual('scripts/bootstrap.py', path) |
46 self.assertEqual(None, revision) | 49 self.assertEqual(None, revision) |
47 self.assertEqual(True, store_last_good) | 50 self.assertEqual(True, store_last_good) |
48 return None, 'foo bar' | 51 return None, 'foo bar' |
49 self.mock(config, 'get_self_config', get_self_config_mock) | 52 self.mock(config, 'get_self_config', get_self_config_mock) |
50 f = bot_code.get_bootstrap('localhost', 'token', None) | 53 f = bot_code.get_bootstrap('localhost', 'token', None) |
51 expected = ( | 54 expected = ( |
52 '#!/usr/bin/env python\n' | 55 '#!/usr/bin/env python\n' |
(...skipping 34 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
87 | 90 |
88 def test_get_bot_version(self): | 91 def test_get_bot_version(self): |
89 actual, additionals = bot_code.get_bot_version('http://localhost') | 92 actual, additionals = bot_code.get_bot_version('http://localhost') |
90 self.assertTrue(re.match(r'^[0-9a-f]{64}$', actual), actual) | 93 self.assertTrue(re.match(r'^[0-9a-f]{64}$', actual), actual) |
91 expected = { | 94 expected = { |
92 'config/bot_config.py': bot_code.get_bot_config().content, | 95 'config/bot_config.py': bot_code.get_bot_config().content, |
93 } | 96 } |
94 self.assertEqual(expected, additionals) | 97 self.assertEqual(expected, additionals) |
95 | 98 |
96 def test_get_swarming_bot_zip(self): | 99 def test_get_swarming_bot_zip(self): |
| 100 local_mc = {'store': {}, 'reads': 0, 'writes': 0} |
| 101 |
| 102 @ndb.tasklet |
| 103 def mock_memcache_get(version, desc, part=None): |
| 104 value = local_mc['store'].get(bot_code.bot_key(version, desc, part)) |
| 105 if value is not None: |
| 106 local_mc['reads'] += 1 |
| 107 raise ndb.Return(value) |
| 108 |
| 109 @ndb.tasklet |
| 110 def mock_memcache_set(value, version, desc, part=None): |
| 111 local_mc['writes'] += 1 |
| 112 key = bot_code.bot_key(version, desc, part) |
| 113 local_mc['store'][key] = value |
| 114 return ndb.Return(None) |
| 115 |
| 116 self.mock(bot_code, 'bot_memcache_set', mock_memcache_set) |
| 117 self.mock(bot_code, 'bot_memcache_get', mock_memcache_get) |
| 118 self.mock(bot_code, 'MAX_MEMCACHED_SIZE_BYTES', 100000) |
| 119 |
| 120 self.assertEqual(0, local_mc['writes']) |
97 zipped_code = bot_code.get_swarming_bot_zip('http://localhost') | 121 zipped_code = bot_code.get_swarming_bot_zip('http://localhost') |
| 122 self.assertEqual(0, local_mc['reads']) |
| 123 self.assertNotEqual(0, local_mc['writes']) |
| 124 |
| 125 # Make sure that we read from memcached if we get it again |
| 126 zipped_code_copy = bot_code.get_swarming_bot_zip('http://localhost') |
| 127 self.assertEqual(local_mc['writes'], local_mc['reads']) |
| 128 # Why not assertEqual? Don't want to dump ~1MB of data if this fails. |
| 129 self.assertTrue(zipped_code == zipped_code_copy) |
| 130 |
98 # Ensure the zip is valid and all the expected files are present. | 131 # Ensure the zip is valid and all the expected files are present. |
99 with zipfile.ZipFile(StringIO.StringIO(zipped_code), 'r') as zip_file: | 132 with zipfile.ZipFile(StringIO.StringIO(zipped_code), 'r') as zip_file: |
100 for i in bot_archive.FILES: | 133 for i in bot_archive.FILES: |
101 with zip_file.open(i) as f: | 134 with zip_file.open(i) as f: |
102 content = f.read() | 135 content = f.read() |
103 if os.path.basename(i) != '__init__.py': | 136 if os.path.basename(i) != '__init__.py': |
104 self.assertTrue(content, i) | 137 self.assertTrue(content, i) |
105 | 138 |
106 temp_dir = tempfile.mkdtemp(prefix='swarming') | 139 temp_dir = tempfile.mkdtemp(prefix='swarming') |
107 try: | 140 try: |
108 # Try running the bot and ensure it can import the required files. (It | 141 # Try running the bot and ensure it can import the required files. (It |
109 # would crash if it failed to import them). | 142 # would crash if it failed to import them). |
110 bot_path = os.path.join(temp_dir, 'swarming_bot.zip') | 143 bot_path = os.path.join(temp_dir, 'swarming_bot.zip') |
111 with open(bot_path, 'wb') as f: | 144 with open(bot_path, 'wb') as f: |
112 f.write(zipped_code) | 145 f.write(zipped_code) |
113 proc = subprocess.Popen( | 146 proc = subprocess.Popen( |
114 [sys.executable, bot_path, 'start_bot', '-h'], | 147 [sys.executable, bot_path, 'start_bot', '-h'], |
115 cwd=temp_dir, | 148 cwd=temp_dir, |
116 stdout=subprocess.PIPE, | 149 stdout=subprocess.PIPE, |
117 stderr=subprocess.STDOUT) | 150 stderr=subprocess.STDOUT) |
118 out = proc.communicate()[0] | 151 out = proc.communicate()[0] |
119 self.assertEqual(0, proc.returncode, out) | 152 self.assertEqual(0, proc.returncode, out) |
120 finally: | 153 finally: |
121 file_path.rmtree(temp_dir) | 154 file_path.rmtree(temp_dir) |
122 | 155 |
| 156 |
123 def test_bootstrap_token(self): | 157 def test_bootstrap_token(self): |
124 tok = bot_code.generate_bootstrap_token() | 158 tok = bot_code.generate_bootstrap_token() |
125 self.assertEqual( | 159 self.assertEqual( |
126 {'for': 'user:joe@localhost'}, bot_code.validate_bootstrap_token(tok)) | 160 {'for': 'user:joe@localhost'}, bot_code.validate_bootstrap_token(tok)) |
127 | 161 |
128 | 162 |
129 if __name__ == '__main__': | 163 if __name__ == '__main__': |
130 fix_encoding.fix_encoding() | 164 fix_encoding.fix_encoding() |
131 logging.basicConfig( | 165 logging.basicConfig( |
132 level=logging.DEBUG if '-v' in sys.argv else logging.ERROR) | 166 level=logging.DEBUG if '-v' in sys.argv else logging.ERROR) |
133 if '-v' in sys.argv: | 167 if '-v' in sys.argv: |
134 unittest.TestCase.maxDiff = None | 168 unittest.TestCase.maxDiff = None |
135 unittest.main() | 169 unittest.main() |
OLD | NEW |