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