Chromium Code Reviews
chromiumcodereview-hr@appspot.gserviceaccount.com (chromiumcodereview-hr) | Please choose your nickname with Settings | Help | Chromium Project | Gerrit Changes | Sign out
(11)

Side by Side Diff: appengine/swarming/server/bot_code_test.py

Issue 2953253003: Replace custom blob gRPC API with ByteStream (Closed)
Patch Set: More responses Created 3 years, 5 months ago
Use n/p to move between diff chunks; N/P to move between comments. Draft comments are only viewable by you.
Jump to:
View unified diff | Download patch
OLDNEW
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
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
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()
OLDNEW

Powered by Google App Engine
This is Rietveld 408576698