blob: 87dc346152a5511947feaca2ef40ef930476ecbb [file] [log] [blame]
xixuanba232a32016-08-25 17:01:59 -07001#!/usr/bin/python
2#
3# Copyright (c) 2016 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7"""Unit tests for frontend/afe/moblab_rpc_interface.py."""
8
9import __builtin__
10# The boto module is only available/used in Moblab for validation of cloud
11# storage access. The module is not available in the test lab environment,
12# and the import error is handled.
xixuanba232a32016-08-25 17:01:59 -070013import ConfigParser
xixuanba232a32016-08-25 17:01:59 -070014import mox
15import StringIO
16import unittest
17
18import common
19
20from autotest_lib.client.common_lib import error
21from autotest_lib.client.common_lib import global_config
22from autotest_lib.client.common_lib import lsbrelease_utils
23from autotest_lib.frontend import setup_django_environment
24from autotest_lib.frontend.afe import frontend_test_utils
25from autotest_lib.frontend.afe import moblab_rpc_interface
26from autotest_lib.frontend.afe import rpc_utils
27from autotest_lib.server import utils
28from autotest_lib.server.hosts import moblab_host
Keith Haddow350d7892017-04-27 15:47:53 -070029from autotest_lib.client.common_lib import utils as common_lib_utils
xixuanba232a32016-08-25 17:01:59 -070030
31
32class MoblabRpcInterfaceTest(mox.MoxTestBase,
33 frontend_test_utils.FrontendTestMixin):
34 """Unit tests for functions in moblab_rpc_interface.py."""
35
36 def setUp(self):
37 super(MoblabRpcInterfaceTest, self).setUp()
38 self._frontend_common_setup(fill_data=False)
39
40
41 def tearDown(self):
42 self._frontend_common_teardown()
43
44
45 def setIsMoblab(self, is_moblab):
46 """Set utils.is_moblab result.
47
48 @param is_moblab: Value to have utils.is_moblab to return.
49 """
50 self.mox.StubOutWithMock(utils, 'is_moblab')
51 utils.is_moblab().AndReturn(is_moblab)
52
53
54 def _mockReadFile(self, path, lines=[]):
55 """Mock out reading a file line by line.
56
57 @param path: Path of the file we are mock reading.
58 @param lines: lines of the mock file that will be returned when
59 readLine() is called.
60 """
61 mockFile = self.mox.CreateMockAnything()
62 for line in lines:
63 mockFile.readline().AndReturn(line)
64 mockFile.readline()
65 mockFile.close()
66 open(path).AndReturn(mockFile)
67
68
69 def testMoblabOnlyDecorator(self):
70 """Ensure the moblab only decorator gates functions properly."""
71 self.setIsMoblab(False)
72 self.mox.ReplayAll()
73 self.assertRaises(error.RPCException,
74 moblab_rpc_interface.get_config_values)
75
76
77 def testGetConfigValues(self):
78 """Ensure that the config object is properly converted to a dict."""
79 self.setIsMoblab(True)
80 config_mock = self.mox.CreateMockAnything()
81 moblab_rpc_interface._CONFIG = config_mock
82 config_mock.get_sections().AndReturn(['section1', 'section2'])
83 config_mock.config = self.mox.CreateMockAnything()
84 config_mock.config.items('section1').AndReturn([('item1', 'value1'),
85 ('item2', 'value2')])
86 config_mock.config.items('section2').AndReturn([('item3', 'value3'),
87 ('item4', 'value4')])
88
89 rpc_utils.prepare_for_serialization(
90 {'section1' : [('item1', 'value1'),
91 ('item2', 'value2')],
92 'section2' : [('item3', 'value3'),
93 ('item4', 'value4')]})
94 self.mox.ReplayAll()
95 moblab_rpc_interface.get_config_values()
96
97
98 def testUpdateConfig(self):
99 """Ensure that updating the config works as expected."""
100 self.setIsMoblab(True)
101 moblab_rpc_interface.os = self.mox.CreateMockAnything()
102
103 self.mox.StubOutWithMock(__builtin__, 'open')
104 self._mockReadFile(global_config.DEFAULT_CONFIG_FILE)
105
106 self.mox.StubOutWithMock(lsbrelease_utils, 'is_moblab')
107 lsbrelease_utils.is_moblab().AndReturn(True)
108
109 self._mockReadFile(global_config.DEFAULT_MOBLAB_FILE,
110 ['[section1]', 'item1: value1'])
111
112 moblab_rpc_interface.os = self.mox.CreateMockAnything()
113 moblab_rpc_interface.os.path = self.mox.CreateMockAnything()
114 moblab_rpc_interface.os.path.exists(
115 moblab_rpc_interface._CONFIG.shadow_file).AndReturn(
116 True)
117 mockShadowFile = self.mox.CreateMockAnything()
118 mockShadowFileContents = StringIO.StringIO()
119 mockShadowFile.__enter__().AndReturn(mockShadowFileContents)
120 mockShadowFile.__exit__(mox.IgnoreArg(), mox.IgnoreArg(),
121 mox.IgnoreArg())
122 open(moblab_rpc_interface._CONFIG.shadow_file,
123 'w').AndReturn(mockShadowFile)
124 moblab_rpc_interface.os.system('sudo reboot')
125
126 self.mox.ReplayAll()
127 moblab_rpc_interface.update_config_handler(
128 {'section1' : [('item1', 'value1'),
129 ('item2', 'value2')],
130 'section2' : [('item3', 'value3'),
131 ('item4', 'value4')]})
132
133 # item1 should not be in the new shadow config as its updated value
134 # matches the original config's value.
135 self.assertEquals(
136 mockShadowFileContents.getvalue(),
137 '[section2]\nitem3 = value3\nitem4 = value4\n\n'
138 '[section1]\nitem2 = value2\n\n')
139
140
141 def testResetConfig(self):
142 """Ensure that reset opens the shadow_config file for writing."""
143 self.setIsMoblab(True)
144 config_mock = self.mox.CreateMockAnything()
145 moblab_rpc_interface._CONFIG = config_mock
146 config_mock.shadow_file = 'shadow_config.ini'
147 self.mox.StubOutWithMock(__builtin__, 'open')
148 mockFile = self.mox.CreateMockAnything()
149 file_contents = self.mox.CreateMockAnything()
150 mockFile.__enter__().AndReturn(file_contents)
151 mockFile.__exit__(mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg())
152 open(config_mock.shadow_file, 'w').AndReturn(mockFile)
153 moblab_rpc_interface.os = self.mox.CreateMockAnything()
154 moblab_rpc_interface.os.system('sudo reboot')
155 self.mox.ReplayAll()
156 moblab_rpc_interface.reset_config_settings()
157
158
xixuanba232a32016-08-25 17:01:59 -0700159 def testSetLaunchControlKey(self):
160 """Ensure that the Launch Control key path supplied is copied correctly.
161 """
162 self.setIsMoblab(True)
163 launch_control_key = '/tmp/launch_control'
164 moblab_rpc_interface.os = self.mox.CreateMockAnything()
165 moblab_rpc_interface.os.path = self.mox.CreateMockAnything()
166 moblab_rpc_interface.os.path.exists(launch_control_key).AndReturn(
167 True)
168 moblab_rpc_interface.shutil = self.mox.CreateMockAnything()
169 moblab_rpc_interface.shutil.copyfile(
170 launch_control_key,
171 moblab_host.MOBLAB_LAUNCH_CONTROL_KEY_LOCATION)
172 moblab_rpc_interface.os.system('sudo restart moblab-devserver-init')
173 self.mox.ReplayAll()
174 moblab_rpc_interface.set_launch_control_key(launch_control_key)
175
176
177 def testGetNetworkInfo(self):
178 """Ensure the network info is properly converted to a dict."""
179 self.setIsMoblab(True)
180
181 self.mox.StubOutWithMock(moblab_rpc_interface, '_get_network_info')
182 moblab_rpc_interface._get_network_info().AndReturn(('10.0.0.1', True))
183 self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization')
184
185 rpc_utils.prepare_for_serialization(
186 {'is_connected': True, 'server_ips': ['10.0.0.1']})
187 self.mox.ReplayAll()
188 moblab_rpc_interface.get_network_info()
189 self.mox.VerifyAll()
190
191
192 def testGetNetworkInfoWithNoIp(self):
193 """Queries network info with no public IP address."""
194 self.setIsMoblab(True)
195
196 self.mox.StubOutWithMock(moblab_rpc_interface, '_get_network_info')
197 moblab_rpc_interface._get_network_info().AndReturn((None, False))
198 self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization')
199
200 rpc_utils.prepare_for_serialization(
201 {'is_connected': False})
202 self.mox.ReplayAll()
203 moblab_rpc_interface.get_network_info()
204 self.mox.VerifyAll()
205
206
207 def testGetNetworkInfoWithNoConnectivity(self):
208 """Queries network info with public IP address but no connectivity."""
209 self.setIsMoblab(True)
210
211 self.mox.StubOutWithMock(moblab_rpc_interface, '_get_network_info')
212 moblab_rpc_interface._get_network_info().AndReturn(('10.0.0.1', False))
213 self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization')
214
215 rpc_utils.prepare_for_serialization(
216 {'is_connected': False, 'server_ips': ['10.0.0.1']})
217 self.mox.ReplayAll()
218 moblab_rpc_interface.get_network_info()
219 self.mox.VerifyAll()
220
221
222 def testGetCloudStorageInfo(self):
223 """Ensure the cloud storage info is properly converted to a dict."""
224 self.setIsMoblab(True)
225 config_mock = self.mox.CreateMockAnything()
226 moblab_rpc_interface._CONFIG = config_mock
227 config_mock.get_config_value(
228 'CROS', 'image_storage_server').AndReturn('gs://bucket1')
229 config_mock.get_config_value(
230 'CROS', 'results_storage_server', default=None).AndReturn(
231 'gs://bucket2')
232 self.mox.StubOutWithMock(moblab_rpc_interface, '_get_boto_config')
233 moblab_rpc_interface._get_boto_config().AndReturn(config_mock)
234 config_mock.sections().AndReturn(['Credentials', 'b'])
235 config_mock.options('Credentials').AndReturn(
236 ['gs_access_key_id', 'gs_secret_access_key'])
237 config_mock.get(
238 'Credentials', 'gs_access_key_id').AndReturn('key')
239 config_mock.get(
240 'Credentials', 'gs_secret_access_key').AndReturn('secret')
241 rpc_utils.prepare_for_serialization(
242 {
243 'gs_access_key_id': 'key',
244 'gs_secret_access_key' : 'secret',
245 'use_existing_boto_file': True,
246 'image_storage_server' : 'gs://bucket1',
247 'results_storage_server' : 'gs://bucket2'
248 })
249 self.mox.ReplayAll()
250 moblab_rpc_interface.get_cloud_storage_info()
251 self.mox.VerifyAll()
252
253
254 def testValidateCloudStorageInfo(self):
255 """ Ensure the cloud storage info validation flow."""
256 self.setIsMoblab(True)
257 cloud_storage_info = {
258 'use_existing_boto_file': False,
259 'gs_access_key_id': 'key',
260 'gs_secret_access_key': 'secret',
261 'image_storage_server': 'gs://bucket1',
262 'results_storage_server': 'gs://bucket2'}
Keith Haddow350d7892017-04-27 15:47:53 -0700263 self.mox.StubOutWithMock(moblab_rpc_interface,
264 '_run_bucket_performance_test')
265 moblab_rpc_interface._run_bucket_performance_test(
266 'key', 'secret', 'gs://bucket1').AndReturn((True, None))
267 rpc_utils.prepare_for_serialization({'status_ok': True })
xixuanba232a32016-08-25 17:01:59 -0700268 self.mox.ReplayAll()
269 moblab_rpc_interface.validate_cloud_storage_info(cloud_storage_info)
270 self.mox.VerifyAll()
271
272
273 def testGetBucketNameFromUrl(self):
274 """Gets bucket name from bucket URL."""
275 self.assertEquals(
276 'bucket_name-123',
277 moblab_rpc_interface._get_bucket_name_from_url(
278 'gs://bucket_name-123'))
279 self.assertEquals(
280 'bucket_name-123',
281 moblab_rpc_interface._get_bucket_name_from_url(
282 'gs://bucket_name-123/'))
283 self.assertEquals(
284 'bucket_name-123',
285 moblab_rpc_interface._get_bucket_name_from_url(
286 'gs://bucket_name-123/a/b/c'))
287 self.assertIsNone(moblab_rpc_interface._get_bucket_name_from_url(
288 'bucket_name-123/a/b/c'))
289
290
xixuanba232a32016-08-25 17:01:59 -0700291 def testGetShadowConfigFromPartialUpdate(self):
292 """Tests getting shadow configuration based on partial upate."""
293 partial_config = {
294 'section1': [
295 ('opt1', 'value1'),
296 ('opt2', 'value2'),
297 ('opt3', 'value3'),
298 ('opt4', 'value4'),
299 ]
300 }
301 shadow_config_str = "[section1]\nopt2 = value2_1\nopt4 = value4_1"
302 shadow_config = ConfigParser.ConfigParser()
303 shadow_config.readfp(StringIO.StringIO(shadow_config_str))
304 original_config = self.mox.CreateMockAnything()
305 self.mox.StubOutWithMock(moblab_rpc_interface, '_read_original_config')
306 self.mox.StubOutWithMock(moblab_rpc_interface, '_read_raw_config')
307 moblab_rpc_interface._read_original_config().AndReturn(original_config)
308 moblab_rpc_interface._read_raw_config(
309 moblab_rpc_interface._CONFIG.shadow_file).AndReturn(shadow_config)
310 original_config.get_config_value(
311 'section1', 'opt1',
312 allow_blank=True, default='').AndReturn('value1')
313 original_config.get_config_value(
314 'section1', 'opt2',
315 allow_blank=True, default='').AndReturn('value2')
316 original_config.get_config_value(
317 'section1', 'opt3',
318 allow_blank=True, default='').AndReturn('blah')
319 original_config.get_config_value(
320 'section1', 'opt4',
321 allow_blank=True, default='').AndReturn('blah')
322 self.mox.ReplayAll()
323 shadow_config = moblab_rpc_interface._get_shadow_config_from_partial_update(
324 partial_config)
325 # opt1 same as the original.
326 self.assertFalse(shadow_config.has_option('section1', 'opt1'))
327 # opt2 reverts back to original
328 self.assertFalse(shadow_config.has_option('section1', 'opt2'))
329 # opt3 is updated from original.
330 self.assertEquals('value3', shadow_config.get('section1', 'opt3'))
331 # opt3 in shadow but updated again.
332 self.assertEquals('value4', shadow_config.get('section1', 'opt4'))
333 self.mox.VerifyAll()
334
335
336 def testGetShadowConfigFromPartialUpdateWithNewSection(self):
337 """
338 Test getting shadown configuration based on partial update with new section.
339 """
340 partial_config = {
341 'section2': [
342 ('opt5', 'value5'),
343 ('opt6', 'value6'),
344 ],
345 }
346 shadow_config_str = "[section1]\nopt2 = value2_1\n"
347 shadow_config = ConfigParser.ConfigParser()
348 shadow_config.readfp(StringIO.StringIO(shadow_config_str))
349 original_config = self.mox.CreateMockAnything()
350 self.mox.StubOutWithMock(moblab_rpc_interface, '_read_original_config')
351 self.mox.StubOutWithMock(moblab_rpc_interface, '_read_raw_config')
352 moblab_rpc_interface._read_original_config().AndReturn(original_config)
353 moblab_rpc_interface._read_raw_config(
354 moblab_rpc_interface._CONFIG.shadow_file).AndReturn(shadow_config)
355 original_config.get_config_value(
356 'section2', 'opt5',
357 allow_blank=True, default='').AndReturn('value5')
358 original_config.get_config_value(
359 'section2', 'opt6',
360 allow_blank=True, default='').AndReturn('blah')
361 self.mox.ReplayAll()
362 shadow_config = moblab_rpc_interface._get_shadow_config_from_partial_update(
363 partial_config)
364 # opt2 is still in shadow
365 self.assertEquals('value2_1', shadow_config.get('section1', 'opt2'))
366 # opt5 is not changed.
367 self.assertFalse(shadow_config.has_option('section2', 'opt5'))
368 # opt6 is updated.
369 self.assertEquals('value6', shadow_config.get('section2', 'opt6'))
370 self.mox.VerifyAll()
371
Keith Haddow3102fd42017-05-10 11:55:29 -0700372 def testGetBuildsForInDirectory(self):
373 config_mock = self.mox.CreateMockAnything()
374 moblab_rpc_interface._CONFIG = config_mock
375 config_mock.get_config_value(
376 'CROS', 'image_storage_server').AndReturn('gs://bucket1/')
377 self.mox.StubOutWithMock(common_lib_utils, 'run')
378 output = self.mox.CreateMockAnything()
379 self.mox.StubOutWithMock(StringIO, 'StringIO', use_mock_anything=True)
380 StringIO.StringIO().AndReturn(output)
381 output.getvalue().AndReturn(
382 """gs://bucket1/dummy/R53-8480.0.0/\ngs://bucket1/dummy/R53-8530.72.0/\n
383 gs://bucket1/dummy/R54-8712.0.0/\ngs://bucket1/dummy/R54-8717.0.0/\n
384 gs://bucket1/dummy/R55-8759.0.0/\n
385 gs://bucket1/dummy/R55-8760.0.0-b5849/\n
386 gs://bucket1/dummy/R56-8995.0.0/\ngs://bucket1/dummy/R56-9001.0.0/\n
387 gs://bucket1/dummy/R57-9202.66.0/\ngs://bucket1/dummy/R58-9331.0.0/\n
388 gs://bucket1/dummy/R58-9334.15.0/\ngs://bucket1/dummy/R58-9334.17.0/\n
389 gs://bucket1/dummy/R58-9334.18.0/\ngs://bucket1/dummy/R58-9334.19.0/\n
390 gs://bucket1/dummy/R58-9334.22.0/\ngs://bucket1/dummy/R58-9334.28.0/\n
391 gs://bucket1/dummy/R58-9334.3.0/\ngs://bucket1/dummy/R58-9334.30.0/\n
392 gs://bucket1/dummy/R58-9334.36.0/\ngs://bucket1/dummy/R58-9334.55.0/\n
393 gs://bucket1/dummy/R58-9334.6.0/\ngs://bucket1/dummy/R58-9334.7.0/\n
394 gs://bucket1/dummy/R58-9334.9.0/\ngs://bucket1/dummy/R59-9346.0.0/\n
395 gs://bucket1/dummy/R59-9372.0.0/\ngs://bucket1/dummy/R59-9387.0.0/\n
396 gs://bucket1/dummy/R59-9436.0.0/\ngs://bucket1/dummy/R59-9452.0.0/\n
397 gs://bucket1/dummy/R59-9453.0.0/\ngs://bucket1/dummy/R59-9455.0.0/\n
398 gs://bucket1/dummy/R59-9460.0.0/\ngs://bucket1/dummy/R59-9460.11.0/\n
399 gs://bucket1/dummy/R59-9460.16.0/\ngs://bucket1/dummy/R59-9460.25.0/\n
400 gs://bucket1/dummy/R59-9460.8.0/\ngs://bucket1/dummy/R59-9460.9.0/\n
401 gs://bucket1/dummy/R60-9472.0.0/\ngs://bucket1/dummy/R60-9491.0.0/\n
402 gs://bucket1/dummy/R60-9492.0.0/\ngs://bucket1/dummy/R60-9497.0.0/\n
403 gs://bucket1/dummy/R60-9500.0.0/""")
404
405 output.close()
406
Keith Haddowc5ec0602017-06-01 16:49:17 -0700407 self.mox.StubOutWithMock(moblab_rpc_interface.GsUtil, 'get_gsutil_cmd')
408 moblab_rpc_interface.GsUtil.get_gsutil_cmd().AndReturn(
409 '/path/to/gsutil')
410
411 common_lib_utils.run('/path/to/gsutil',
Keith Haddow3102fd42017-05-10 11:55:29 -0700412 args=('ls', 'gs://bucket1/dummy'),
413 stdout_tee=mox.IgnoreArg()).AndReturn(output)
414 self.mox.ReplayAll()
415 expected_results = ['dummy/R60-9500.0.0', 'dummy/R60-9497.0.0',
416 'dummy/R60-9492.0.0', 'dummy/R60-9491.0.0', 'dummy/R60-9472.0.0',
417 'dummy/R59-9460.25.0', 'dummy/R59-9460.16.0', 'dummy/R59-9460.11.0',
418 'dummy/R59-9460.9.0', 'dummy/R59-9460.8.0', 'dummy/R58-9334.55.0',
419 'dummy/R58-9334.36.0', 'dummy/R58-9334.30.0', 'dummy/R58-9334.28.0',
420 'dummy/R58-9334.22.0']
421 actual_results = moblab_rpc_interface._get_builds_for_in_directory(
422 "dummy",3, 5)
423 self.assertEquals(expected_results, actual_results)
424 self.mox.VerifyAll()
xixuanba232a32016-08-25 17:01:59 -0700425
Keith Haddow350d7892017-04-27 15:47:53 -0700426 def testRunBucketPerformanceTestFail(self):
Keith Haddowc5ec0602017-06-01 16:49:17 -0700427 self.mox.StubOutWithMock(moblab_rpc_interface.GsUtil, 'get_gsutil_cmd')
428 moblab_rpc_interface.GsUtil.get_gsutil_cmd().AndReturn(
429 '/path/to/gsutil')
Keith Haddow350d7892017-04-27 15:47:53 -0700430 self.mox.StubOutWithMock(common_lib_utils, 'run')
Keith Haddowc5ec0602017-06-01 16:49:17 -0700431 common_lib_utils.run('/path/to/gsutil',
Keith Haddow350d7892017-04-27 15:47:53 -0700432 args=(
433 '-o', 'Credentials:gs_access_key_id=key',
434 '-o', 'Credentials:gs_secret_access_key=secret',
435 'perfdiag', '-s', '1K',
436 '-o', 'testoutput',
437 '-n', '10',
438 'gs://bucket1')).AndRaise(
439 error.CmdError("fakecommand", common_lib_utils.CmdResult(),
440 "xxxxxx<Error>yyyyyyyyyy</Error>"))
441
442 self.mox.ReplayAll()
443 self.assertRaisesRegexp(
444 moblab_rpc_interface.BucketPerformanceTestException,
445 '<Error>yyyyyyyyyy',
446 moblab_rpc_interface._run_bucket_performance_test,
447 'key', 'secret', 'gs://bucket1', '1K', '10', 'testoutput')
448 self.mox.VerifyAll()
449
Keith Haddow21e33a52017-05-17 15:42:58 -0700450 def testEnableNotificationUsingCredentialsInBucketFail(self):
451 config_mock = self.mox.CreateMockAnything()
452 moblab_rpc_interface._CONFIG = config_mock
453 config_mock.get_config_value(
454 'CROS', 'image_storage_server').AndReturn('gs://bucket1/')
Keith Haddowc5ec0602017-06-01 16:49:17 -0700455
456 self.mox.StubOutWithMock(moblab_rpc_interface.GsUtil, 'get_gsutil_cmd')
457 moblab_rpc_interface.GsUtil.get_gsutil_cmd().AndReturn(
458 '/path/to/gsutil')
459
Keith Haddow21e33a52017-05-17 15:42:58 -0700460 self.mox.StubOutWithMock(common_lib_utils, 'run')
Keith Haddowc5ec0602017-06-01 16:49:17 -0700461 common_lib_utils.run('/path/to/gsutil',
Keith Haddow21e33a52017-05-17 15:42:58 -0700462 args=('cp', 'gs://bucket1/pubsub-key-do-not-delete.json',
463 '/tmp')).AndRaise(
464 error.CmdError("fakecommand", common_lib_utils.CmdResult(), ""))
465 self.mox.ReplayAll()
466 moblab_rpc_interface._enable_notification_using_credentials_in_bucket()
467
468 def testEnableNotificationUsingCredentialsInBucketSuccess(self):
469 config_mock = self.mox.CreateMockAnything()
470 moblab_rpc_interface._CONFIG = config_mock
471 config_mock.get_config_value(
472 'CROS', 'image_storage_server').AndReturn('gs://bucket1/')
Keith Haddowc5ec0602017-06-01 16:49:17 -0700473
474 self.mox.StubOutWithMock(moblab_rpc_interface.GsUtil, 'get_gsutil_cmd')
475 moblab_rpc_interface.GsUtil.get_gsutil_cmd().AndReturn(
476 '/path/to/gsutil')
477
Keith Haddow21e33a52017-05-17 15:42:58 -0700478 self.mox.StubOutWithMock(common_lib_utils, 'run')
Keith Haddowc5ec0602017-06-01 16:49:17 -0700479 common_lib_utils.run('/path/to/gsutil',
Keith Haddow21e33a52017-05-17 15:42:58 -0700480 args=('cp', 'gs://bucket1/pubsub-key-do-not-delete.json',
481 '/tmp'))
482 moblab_rpc_interface.shutil = self.mox.CreateMockAnything()
483 moblab_rpc_interface.shutil.copyfile(
484 '/tmp/pubsub-key-do-not-delete.json',
485 moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION)
486 self.mox.StubOutWithMock(moblab_rpc_interface, '_update_partial_config')
487 moblab_rpc_interface._update_partial_config(
488 {'CROS': [(moblab_rpc_interface._CLOUD_NOTIFICATION_ENABLED, True)]}
489 )
490 self.mox.ReplayAll()
491 moblab_rpc_interface._enable_notification_using_credentials_in_bucket()
492
Matt Mallettb0f8dc72018-01-25 15:15:58 -0800493 def testInstallSystemUpdate(self):
494 update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT
495
496 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'check_call')
497 moblab_rpc_interface.subprocess.check_call(['sudo',
498 update_engine_client, '--update'])
Matt Mallettb0f8dc72018-01-25 15:15:58 -0800499 moblab_rpc_interface.subprocess.check_call(['sudo',
Matt Mallettc2235052018-01-30 16:47:08 -0800500 update_engine_client, '--is_reboot_needed'])
Matt Mallettb0f8dc72018-01-25 15:15:58 -0800501
502 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'call')
503 moblab_rpc_interface.subprocess.call(['sudo', update_engine_client,
504 '--reboot'])
505
506 self.mox.ReplayAll()
507 moblab_rpc_interface._install_system_update()
508
509
510 def testGetSystemUpdateStatus(self):
511 update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT
512 update_status = ('LAST_CHECKED_TIME=1516753795\n'
513 'PROGRESS=0.220121\n'
514 'CURRENT_OP=UPDATE_STATUS_DOWNLOADING\n'
515 'NEW_VERSION=10032.89.0\n'
516 'NEW_SIZE=782805733')
517
518 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess,
519 'check_output')
520 moblab_rpc_interface.subprocess.check_output(['sudo',
521 update_engine_client, '--status']).AndReturn(
522 update_status)
523
524 self.mox.ReplayAll()
525 output = moblab_rpc_interface._get_system_update_status()
526
527 self.assertEquals(output['PROGRESS'], '0.220121')
528 self.assertEquals(output['CURRENT_OP'], 'UPDATE_STATUS_DOWNLOADING')
529 self.assertEquals(output['NEW_VERSION'], '10032.89.0')
530 self.assertEquals(output['NEW_SIZE'], '782805733')
531
532 def testCheckForSystemUpdate(self):
533 update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT
534
535 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'call')
536 moblab_rpc_interface.subprocess.call(['sudo', update_engine_client,
537 '--check_for_update'])
538
539 self.mox.ReplayAll()
540 moblab_rpc_interface._check_for_system_update()
541
Matt Mallett31c1eec2018-01-30 16:13:26 -0800542 def testGetConnectedDutBoardModels(self):
543 # setting up mocks for 2 duts with different boards and models
544 mock_minnie_labels = [
545 self.mox.CreateMockAnything(),
546 self.mox.CreateMockAnything(),
547 ]
548 mock_minnie_labels[0].name = 'board:veyron_minnie'
549 mock_minnie_labels[1].name = 'model:veyron_minnie'
550 mock_minnie = self.mox.CreateMockAnything()
551 mock_minnie.label_list = mock_minnie_labels
552
553 mock_bruce_labels = [
554 self.mox.CreateMockAnything(),
555 self.mox.CreateMockAnything()
556 ]
557 mock_bruce_labels[0].name = 'board:carl'
558 mock_bruce_labels[1].name = 'model:bruce'
559 mock_bruce = self.mox.CreateMockAnything()
560 mock_bruce.label_list = mock_bruce_labels
561 hosts = [mock_minnie, mock_bruce]
562
563 # stub out the host query calls
564 self.mox.StubOutWithMock(moblab_rpc_interface.rpc_utils,
565 'get_host_query')
566 moblab_rpc_interface.rpc_utils.get_host_query(
567 (), False, True, {}).AndReturn(hosts)
568
569 self.mox.StubOutWithMock(moblab_rpc_interface.models.Host.objects,
570 'populate_relationships'),
571 moblab_rpc_interface.models.Host.objects.populate_relationships(hosts,
572 moblab_rpc_interface.models.Label, 'label_list')
573
574 expected = [{
575 'model': 'bruce',
576 'board': 'carl'
577 },
578 {
579 'model': 'veyron_minnie',
580 'board': 'veyron_minnie'
581 }]
582
583 self.mox.ReplayAll()
584 output = moblab_rpc_interface._get_connected_dut_board_models()
585 self.assertEquals(output, expected)
586 # test sorting
587 self.assertEquals(output[0]['model'], 'bruce')
588
Keith Haddow350d7892017-04-27 15:47:53 -0700589
xixuanba232a32016-08-25 17:01:59 -0700590if __name__ == '__main__':
591 unittest.main()