OLD | NEW |
---|---|
(Empty) | |
1 #!/usr/bin/env python | |
2 # Copyright 2016 The LUCI Authors. All rights reserved. | |
3 # Use of this source code is governed under the Apache License, Version 2.0 | |
4 # that can be found in the LICENSE file. | |
5 | |
6 import logging | |
7 import sys | |
8 import threading | |
9 import time | |
10 import unittest | |
11 | |
12 import test_env_bot_code | |
13 test_env_bot_code.setup_test_env() | |
14 | |
15 from depot_tools import auto_stub | |
16 | |
17 try: | |
18 import remote_client_grpc | |
19 except ImportError as e: | |
20 print('Could not import gRPC remote client, likely due to missing grpc ' | |
21 'library. Skipping tests.') | |
22 sys.exit(0) | |
23 | |
24 class FakeBotServiceStub(object): | |
25 def __init__(self, testobj): | |
26 self._testobj = testobj | |
27 | |
28 def TaskUpdate(self, request, **_kwargs): | |
29 return self._testobj._handle_call('TaskUpdate', request) | |
30 | |
31 def Handshake(self, request, **_kwargs): | |
32 return self._testobj._handle_call('Handshake', request) | |
33 | |
34 def Poll(self, request, **_kwargs): | |
35 return self._testobj._handle_call('Poll', request) | |
36 | |
37 | |
38 | |
M-A Ruel
2016/12/14 21:21:52
inconsistent spacing between file level symbols
aludwin
2016/12/14 22:10:17
Done.
| |
39 | |
40 class TestRemoteClientGrpc(auto_stub.TestCase): | |
41 def setUp(self): | |
42 super(TestRemoteClientGrpc, self).setUp() | |
43 self._client = remote_client_grpc.RemoteClientGrpc('1.2.3.4:90') | |
44 self._client._stub = FakeBotServiceStub(self) | |
45 self._expected = [] | |
46 # self.slept = 0 | |
M-A Ruel
2016/12/14 21:21:52
remove before committing
aludwin
2016/12/14 22:10:17
Done.
| |
47 # def sleep_mock(t): | |
48 # self.slept += t | |
49 # self.mock(time, 'sleep', sleep_mock) | |
50 | |
51 def _handle_call(self, method, request): | |
52 # Pop off the first item on the list | |
53 self.assertTrue(len(self._expected) > 0) | |
54 expected, self._expected = self._expected[0], self._expected[1:] | |
55 # Each element of the "expected" array should be a 3-tuple: | |
56 # * The name of the method (eg 'TaskUpdate') | |
57 # * The proto request | |
58 # * The proto response OR a callable that accepts the request | |
59 self.assertEqual(method, expected[0]) | |
60 self.assertEqual(request, expected[1]) | |
61 response = expected[2] | |
62 if callable(response): | |
63 return response(request) | |
64 return response | |
65 | |
66 def get_bot_attributes_dict(self): | |
67 """Gets the attributes of a basic bot""" | |
68 return { | |
69 'version': '123', | |
70 'dimensions': { | |
71 'pool': ['default'], | |
M-A Ruel
2016/12/14 21:21:52
sort keys
aludwin
2016/12/14 22:10:17
Done.
| |
72 'mammal': ['ferrett', 'wombat'], | |
73 }, | |
74 'state': { | |
75 'audio': ['Rachmaninov', 'Stravinsky'], | |
76 'cost_usd_hour': 3.141, | |
77 'cpu': 'Intel 8080', | |
78 'sleep_streak': 8, | |
79 'disks': { | |
80 'mount': 'path/to/mount', | |
81 'volume': 'turned-to-eleven', | |
82 }, | |
83 }, | |
84 } | |
85 | |
86 def get_bot_attributes_proto(self): | |
87 """Gets the attributes proto of a basic bot; must match prior function""" | |
88 attributes = remote_client_grpc.swarming_bot_pb2.Attributes() | |
89 attributes.version = '123' | |
90 # Note that dimensions are sorted | |
91 d1 = attributes.dimensions.add() | |
92 d1.name = 'mammal' | |
93 d1.values.extend(['ferrett', 'wombat']) | |
94 d2 = attributes.dimensions.add() | |
95 d2.name = 'pool' | |
96 d2.values.append('default') | |
97 attributes.state.audio.extend(['Rachmaninov', 'Stravinsky']) | |
98 attributes.state.cost_usd_hour = 3.141 | |
99 attributes.state.cpu = 'Intel 8080' | |
100 attributes.state.sleep_streak = 8 | |
101 disk1 = attributes.state.disks.fields['mount'] | |
102 disk1.string_value = 'path/to/mount' | |
103 disk2 = attributes.state.disks.fields['volume'] | |
104 disk2.string_value = 'turned-to-eleven' | |
105 return attributes | |
106 | |
107 def test_handshake(self): | |
108 """Tests the handshake proto""" | |
109 msg_req = remote_client_grpc.swarming_bot_pb2.HandshakeRequest() | |
110 msg_req.attributes.CopyFrom(self.get_bot_attributes_proto()) | |
111 | |
112 # Create proto response | |
113 msg_rsp = remote_client_grpc.swarming_bot_pb2.HandshakeResponse() | |
114 msg_rsp.server_version = '101' | |
115 msg_rsp.bot_version = '102' | |
116 d1 = msg_rsp.bot_group_cfg.dimensions.add() | |
117 d1.name = 'mammal' | |
118 d1.values.extend(['kangaroo', 'emu']) | |
119 | |
120 # Execute call and verify response | |
121 expected_call = ('Handshake', msg_req, msg_rsp) | |
122 self._expected.append(expected_call) | |
123 response = self._client.do_handshake(self.get_bot_attributes_dict()) | |
124 self.assertEqual(response, { | |
125 'server_version': u'101', | |
126 'bot_version': u'102', | |
127 'bot_group_cfg_version': u'', | |
128 'bot_group_cfg': { | |
129 'dimensions': { | |
130 u'mammal': [u'kangaroo', u'emu'], | |
131 }, | |
132 }, | |
133 }) | |
134 | |
135 def test_poll_manifest(self): | |
136 """Verifies that we can generate a reasonable manifest from a proto""" | |
137 msg_req = remote_client_grpc.swarming_bot_pb2.PollRequest() | |
138 msg_req.attributes.CopyFrom(self.get_bot_attributes_proto()) | |
139 | |
140 # Create dict response | |
141 dict_rsp = { | |
142 'bot_id': u'baby_bot', | |
143 'dimensions': { | |
144 u'mammal': u'emu', | |
145 u'pool': u'dead', | |
146 }, | |
147 'env': { | |
148 u'co2': u'400ppm', | |
149 u'n2': u'80%', | |
150 u'o2': u'20%', | |
151 }, | |
152 'grace_period': 1, | |
153 'hard_timeout': 2, | |
154 'io_timeout': 3, | |
155 'isolated': { | |
156 'namespace': u'default-gzip', | |
157 'input': u'123abc', | |
158 'server': '1.2.3.4:90', # from when client was created | |
159 }, | |
160 'outputs': [u'foo.o', u'foo.log'], | |
161 'task_id': u'ifyouchoosetoacceptit', | |
162 'command': None, | |
163 'extra_args': None, | |
164 } | |
165 # Create proto response | |
166 msg_rsp = remote_client_grpc.swarming_bot_pb2.PollResponse() | |
167 msg_rsp.cmd = remote_client_grpc.swarming_bot_pb2.PollResponse.RUN | |
168 msg_rsp.manifest.bot_id = 'baby_bot' | |
169 msg_rsp.manifest.dimensions['mammal'] = 'emu' | |
170 msg_rsp.manifest.dimensions['pool'] = 'dead' | |
171 msg_rsp.manifest.env['co2'] = '400ppm' | |
172 msg_rsp.manifest.env['n2'] = '80%' | |
173 msg_rsp.manifest.env['o2'] = '20%' | |
174 msg_rsp.manifest.grace_period = 1 | |
175 msg_rsp.manifest.hard_timeout = 2 | |
176 msg_rsp.manifest.io_timeout = 3 | |
177 msg_rsp.manifest.isolated.namespace = 'default-gzip' | |
178 msg_rsp.manifest.isolated.input = '123abc' | |
179 msg_rsp.manifest.outputs.extend(['foo.o', 'foo.log']) | |
180 msg_rsp.manifest.task_id = 'ifyouchoosetoacceptit' | |
181 | |
182 # Execute call and verify response | |
183 expected_call = ('Poll', msg_req, msg_rsp) | |
184 self._expected.append(expected_call) | |
185 cmd, value = self._client.poll(self.get_bot_attributes_dict()) | |
186 self.assertEqual(cmd, 'run') | |
187 self.assertEqual(value, dict_rsp) | |
188 | |
189 def test_update_no_output_or_code(self): | |
190 """Includes only a basic param set""" | |
191 # Create params dict | |
192 params = { | |
193 'bot_overhead': 3.141, | |
194 'cost_usd': 0.001, | |
195 } | |
196 | |
197 # Create expected request proto | |
198 msg_req = remote_client_grpc.swarming_bot_pb2.TaskUpdateRequest() | |
199 msg_req.id = 'baby_bot' | |
200 msg_req.task_id = 'abc123' | |
201 msg_req.bot_overhead = 3.141 | |
202 msg_req.cost_usd = 0.001 | |
203 | |
204 # Create proto response | |
205 msg_rsp = remote_client_grpc.swarming_bot_pb2.TaskUpdateResponse() | |
206 | |
207 # Execute call and verify response | |
208 expected_call = ('TaskUpdate', msg_req, msg_rsp) | |
209 self._expected.append(expected_call) | |
210 response = self._client.post_task_update('abc123', 'baby_bot', params) | |
211 self.assertTrue(response) | |
212 | |
213 def test_update_final(self): | |
214 """Includes output chunk, exit code and full param set""" | |
215 # Create params dict | |
216 params = { | |
217 'outputs_ref': { | |
218 'isolated': 'def987', | |
219 'isolatedserver': 'foo', | |
220 'namespace': 'default-gzip', | |
221 }, | |
222 'isolated_stats': { | |
223 'upload': { | |
224 'duration': 0.5, | |
225 'items_cold': 'YWJj', # b64(u'abc') | |
226 }, | |
227 'download': { | |
228 'initial_size': 1234567890, | |
229 }, | |
230 }, | |
231 } | |
232 | |
233 # Create expected request proto | |
234 msg_req = remote_client_grpc.swarming_bot_pb2.TaskUpdateRequest() | |
235 msg_req.id = 'baby_bot' | |
236 msg_req.task_id = 'abc123' | |
237 msg_req.exit_status.code = -5 | |
238 msg_req.output_chunk.data = 'abc' | |
239 msg_req.output_chunk.offset = 7 | |
240 msg_req.outputs_ref.isolated = 'def987' | |
241 msg_req.outputs_ref.isolatedserver = 'foo' | |
242 msg_req.outputs_ref.namespace = 'default-gzip' | |
243 msg_req.isolated_stats.upload.duration = 0.5 | |
244 msg_req.isolated_stats.upload.items_cold = 'abc' | |
245 msg_req.isolated_stats.download.initial_size = 1234567890 | |
246 | |
247 # Create proto response | |
248 msg_rsp = remote_client_grpc.swarming_bot_pb2.TaskUpdateResponse() | |
249 | |
250 # Execute call and verify response | |
251 expected_call = ('TaskUpdate', msg_req, msg_rsp) | |
252 self._expected.append(expected_call) | |
253 response = self._client.post_task_update('abc123', 'baby_bot', params, | |
254 ['abc', 7], -5) | |
255 self.assertTrue(response) | |
256 | |
257 if __name__ == '__main__': | |
258 logging.basicConfig( | |
259 level=logging.DEBUG if '-v' in sys.argv else logging.CRITICAL) | |
260 unittest.TestCase.maxDiff = None | |
261 unittest.main() | |
OLD | NEW |