1 #!/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 9 import __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. 13 import ConfigParser 14 import mox 15 import StringIO 16 import unittest 17 18 import common 19 20 from autotest_lib.client.common_lib import error 21 from autotest_lib.client.common_lib import global_config 22 from autotest_lib.client.common_lib import lsbrelease_utils 23 from autotest_lib.frontend import setup_django_environment 24 from autotest_lib.frontend.afe import frontend_test_utils 25 from autotest_lib.frontend.afe import moblab_rpc_interface 26 from autotest_lib.frontend.afe import rpc_utils 27 from autotest_lib.server import utils 28 from autotest_lib.server.hosts import moblab_host 29 from autotest_lib.client.common_lib import utils as common_lib_utils 30 31 32 class 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 159 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'} 263 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 }) 268 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 291 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 372 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 407 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', 412 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() 425 426 def testRunBucketPerformanceTestFail(self): 427 self.mox.StubOutWithMock(moblab_rpc_interface.GsUtil, 'get_gsutil_cmd') 428 moblab_rpc_interface.GsUtil.get_gsutil_cmd().AndReturn( 429 '/path/to/gsutil') 430 self.mox.StubOutWithMock(common_lib_utils, 'run') 431 common_lib_utils.run('/path/to/gsutil', 432 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 450 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/') 455 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 460 self.mox.StubOutWithMock(common_lib_utils, 'run') 461 common_lib_utils.run('/path/to/gsutil', 462 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/') 473 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 478 self.mox.StubOutWithMock(common_lib_utils, 'run') 479 common_lib_utils.run('/path/to/gsutil', 480 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 493 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']) 499 moblab_rpc_interface.subprocess.check_call(['sudo', 500 update_engine_client, '--is_reboot_needed']) 501 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 def testInstallSystemUpdateError(self): 510 update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT 511 512 error_message = ('ERROR_CODE=37\n' 513 'ERROR_MESSAGE=ErrorCode::kOmahaErrorInHTTPResponse') 514 515 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'check_call') 516 moblab_rpc_interface.subprocess.check_call(['sudo', 517 update_engine_client, '--update']).AndRaise( 518 moblab_rpc_interface.subprocess.CalledProcessError(1, 519 'sudo')) 520 521 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 522 'check_output') 523 moblab_rpc_interface.subprocess.check_output(['sudo', 524 update_engine_client, '--last_attempt_error']).AndReturn( 525 error_message) 526 527 self.mox.ReplayAll() 528 try: 529 moblab_rpc_interface._install_system_update() 530 except moblab_rpc_interface.error.RPCException as e: 531 self.assertEquals(str(e), error_message) 532 533 534 def testGetSystemUpdateStatus(self): 535 update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT 536 update_status = ('LAST_CHECKED_TIME=1516753795\n' 537 'PROGRESS=0.220121\n' 538 'CURRENT_OP=UPDATE_STATUS_DOWNLOADING\n' 539 'NEW_VERSION=10032.89.0\n' 540 'NEW_SIZE=782805733') 541 542 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 543 'check_output') 544 moblab_rpc_interface.subprocess.check_output(['sudo', 545 update_engine_client, '--status']).AndReturn( 546 update_status) 547 548 self.mox.ReplayAll() 549 output = moblab_rpc_interface._get_system_update_status() 550 551 self.assertEquals(output['PROGRESS'], '0.220121') 552 self.assertEquals(output['CURRENT_OP'], 'UPDATE_STATUS_DOWNLOADING') 553 self.assertEquals(output['NEW_VERSION'], '10032.89.0') 554 self.assertEquals(output['NEW_SIZE'], '782805733') 555 556 def testCheckForSystemUpdate(self): 557 update_engine_client = moblab_rpc_interface._UPDATE_ENGINE_CLIENT 558 559 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'call') 560 moblab_rpc_interface.subprocess.call(['sudo', update_engine_client, 561 '--check_for_update']) 562 563 self.mox.StubOutWithMock(moblab_rpc_interface, 564 '_get_system_update_status') 565 for i in range(0,4): 566 moblab_rpc_interface._get_system_update_status().AndReturn( 567 dict({'CURRENT_OP': 'UPDATE_STATUS_CHECKING_FOR_UPDATE'}) 568 ) 569 moblab_rpc_interface._get_system_update_status().AndReturn( 570 dict({'CURRENT_OP': 'UPDATE_STATUS_DOWNLOADING'}) 571 ) 572 self.mox.ReplayAll() 573 moblab_rpc_interface._check_for_system_update() 574 575 def testGetConnectedDutBoardModels(self): 576 # setting up mocks for 2 duts with different boards and models 577 mock_minnie_labels = [ 578 self.mox.CreateMockAnything(), 579 self.mox.CreateMockAnything(), 580 ] 581 mock_minnie_labels[0].name = 'board:veyron_minnie' 582 mock_minnie_labels[1].name = 'model:veyron_minnie' 583 mock_minnie = self.mox.CreateMockAnything() 584 mock_minnie.label_list = mock_minnie_labels 585 586 mock_bruce_labels = [ 587 self.mox.CreateMockAnything(), 588 self.mox.CreateMockAnything() 589 ] 590 mock_bruce_labels[0].name = 'board:carl' 591 mock_bruce_labels[1].name = 'model:bruce' 592 mock_bruce = self.mox.CreateMockAnything() 593 mock_bruce.label_list = mock_bruce_labels 594 hosts = [mock_minnie, mock_bruce] 595 596 # stub out the host query calls 597 self.mox.StubOutWithMock(moblab_rpc_interface.rpc_utils, 598 'get_host_query') 599 moblab_rpc_interface.rpc_utils.get_host_query( 600 (), False, True, {}).AndReturn(hosts) 601 602 self.mox.StubOutWithMock(moblab_rpc_interface.models.Host.objects, 603 'populate_relationships'), 604 moblab_rpc_interface.models.Host.objects.populate_relationships(hosts, 605 moblab_rpc_interface.models.Label, 'label_list') 606 607 expected = [{ 608 'model': 'bruce', 609 'board': 'carl' 610 }, 611 { 612 'model': 'veyron_minnie', 613 'board': 'veyron_minnie' 614 }] 615 616 self.mox.ReplayAll() 617 output = moblab_rpc_interface._get_connected_dut_board_models() 618 self.assertEquals(output, expected) 619 # test sorting 620 self.assertEquals(output[0]['model'], 'bruce') 621 622 def testAllDutConnections(self): 623 leases = { 624 '192.168.0.20': '3c:52:82:5f:15:20', 625 '192.168.0.30': '3c:52:82:5f:15:21' 626 } 627 628 # stub out all of the multiprocessing 629 mock_value = self.mox.CreateMockAnything() 630 mock_value.value = True 631 mock_process = self.mox.CreateMockAnything() 632 633 for key in leases: 634 mock_process.start() 635 for key in leases: 636 mock_process.join() 637 638 self.mox.StubOutWithMock( 639 moblab_rpc_interface, 'multiprocessing') 640 641 for key in leases: 642 moblab_rpc_interface.multiprocessing.Value( 643 mox.IgnoreArg()).AndReturn(mock_value) 644 moblab_rpc_interface.multiprocessing.Process( 645 target=mox.IgnoreArg(), args=mox.IgnoreArg()).AndReturn( 646 mock_process) 647 648 self.mox.ReplayAll() 649 650 expected = { 651 '192.168.0.20': { 652 'mac_address': '3c:52:82:5f:15:20', 653 'ssh_connection_ok': True 654 }, 655 '192.168.0.30': { 656 'mac_address': '3c:52:82:5f:15:21', 657 'ssh_connection_ok': True 658 } 659 } 660 661 connected_duts = moblab_rpc_interface._test_all_dut_connections(leases) 662 self.assertDictEqual(expected, connected_duts) 663 664 def testAllDutConnectionsFailure(self): 665 leases = { 666 '192.168.0.20': '3c:52:82:5f:15:20', 667 '192.168.0.30': '3c:52:82:5f:15:21' 668 } 669 670 # stub out all of the multiprocessing 671 mock_value = self.mox.CreateMockAnything() 672 mock_value.value = False 673 mock_process = self.mox.CreateMockAnything() 674 675 for key in leases: 676 mock_process.start() 677 for key in leases: 678 mock_process.join() 679 680 self.mox.StubOutWithMock( 681 moblab_rpc_interface, 'multiprocessing') 682 683 for key in leases: 684 moblab_rpc_interface.multiprocessing.Value( 685 mox.IgnoreArg()).AndReturn(mock_value) 686 moblab_rpc_interface.multiprocessing.Process( 687 target=mox.IgnoreArg(), args=mox.IgnoreArg()).AndReturn( 688 mock_process) 689 690 self.mox.ReplayAll() 691 692 expected = { 693 '192.168.0.20': { 694 'mac_address': '3c:52:82:5f:15:20', 695 'ssh_connection_ok': False 696 }, 697 '192.168.0.30': { 698 'mac_address': '3c:52:82:5f:15:21', 699 'ssh_connection_ok': False 700 } 701 } 702 703 connected_duts = moblab_rpc_interface._test_all_dut_connections(leases) 704 self.assertDictEqual(expected, connected_duts) 705 706 def testDutSshConnection(self): 707 good_ip = '192.168.0.20' 708 bad_ip = '192.168.0.30' 709 cmd = ('ssh -o ConnectTimeout=2 -o StrictHostKeyChecking=no ' 710 "root@%s 'timeout 2 cat /etc/lsb-release'") 711 712 self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 713 'check_output') 714 moblab_rpc_interface.subprocess.check_output( 715 cmd % good_ip, shell=True).AndReturn('CHROMEOS_RELEASE_APPID') 716 717 moblab_rpc_interface.subprocess.check_output( 718 cmd % bad_ip, shell=True).AndRaise( 719 moblab_rpc_interface.subprocess.CalledProcessError(1, cmd)) 720 721 self.mox.ReplayAll() 722 self.assertEquals( 723 moblab_rpc_interface._test_dut_ssh_connection(good_ip), True) 724 self.assertEquals( 725 moblab_rpc_interface._test_dut_ssh_connection(bad_ip), False) 726 727 728 if __name__ == '__main__': 729 unittest.main() 730
