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 try: 14 import boto 15 except ImportError: 16 boto = None 17 import ConfigParser 18 import logging 19 import mox 20 import StringIO 21 import unittest 22 23 import common 24 25 from autotest_lib.client.common_lib import error 26 from autotest_lib.client.common_lib import global_config 27 from autotest_lib.client.common_lib import lsbrelease_utils 28 from autotest_lib.frontend import setup_django_environment 29 from autotest_lib.frontend.afe import frontend_test_utils 30 from autotest_lib.frontend.afe import moblab_rpc_interface 31 from autotest_lib.frontend.afe import rpc_utils 32 from autotest_lib.server import utils 33 from autotest_lib.server.hosts import moblab_host 34 35 36 class MoblabRpcInterfaceTest(mox.MoxTestBase, 37 frontend_test_utils.FrontendTestMixin): 38 """Unit tests for functions in moblab_rpc_interface.py.""" 39 40 def setUp(self): 41 super(MoblabRpcInterfaceTest, self).setUp() 42 self._frontend_common_setup(fill_data=False) 43 44 45 def tearDown(self): 46 self._frontend_common_teardown() 47 48 49 def setIsMoblab(self, is_moblab): 50 """Set utils.is_moblab result. 51 52 @param is_moblab: Value to have utils.is_moblab to return. 53 """ 54 self.mox.StubOutWithMock(utils, 'is_moblab') 55 utils.is_moblab().AndReturn(is_moblab) 56 57 58 def _mockReadFile(self, path, lines=[]): 59 """Mock out reading a file line by line. 60 61 @param path: Path of the file we are mock reading. 62 @param lines: lines of the mock file that will be returned when 63 readLine() is called. 64 """ 65 mockFile = self.mox.CreateMockAnything() 66 for line in lines: 67 mockFile.readline().AndReturn(line) 68 mockFile.readline() 69 mockFile.close() 70 open(path).AndReturn(mockFile) 71 72 73 def testMoblabOnlyDecorator(self): 74 """Ensure the moblab only decorator gates functions properly.""" 75 self.setIsMoblab(False) 76 self.mox.ReplayAll() 77 self.assertRaises(error.RPCException, 78 moblab_rpc_interface.get_config_values) 79 80 81 def testGetConfigValues(self): 82 """Ensure that the config object is properly converted to a dict.""" 83 self.setIsMoblab(True) 84 config_mock = self.mox.CreateMockAnything() 85 moblab_rpc_interface._CONFIG = config_mock 86 config_mock.get_sections().AndReturn(['section1', 'section2']) 87 config_mock.config = self.mox.CreateMockAnything() 88 config_mock.config.items('section1').AndReturn([('item1', 'value1'), 89 ('item2', 'value2')]) 90 config_mock.config.items('section2').AndReturn([('item3', 'value3'), 91 ('item4', 'value4')]) 92 93 rpc_utils.prepare_for_serialization( 94 {'section1' : [('item1', 'value1'), 95 ('item2', 'value2')], 96 'section2' : [('item3', 'value3'), 97 ('item4', 'value4')]}) 98 self.mox.ReplayAll() 99 moblab_rpc_interface.get_config_values() 100 101 102 def testUpdateConfig(self): 103 """Ensure that updating the config works as expected.""" 104 self.setIsMoblab(True) 105 moblab_rpc_interface.os = self.mox.CreateMockAnything() 106 107 self.mox.StubOutWithMock(__builtin__, 'open') 108 self._mockReadFile(global_config.DEFAULT_CONFIG_FILE) 109 110 self.mox.StubOutWithMock(lsbrelease_utils, 'is_moblab') 111 lsbrelease_utils.is_moblab().AndReturn(True) 112 113 self._mockReadFile(global_config.DEFAULT_MOBLAB_FILE, 114 ['[section1]', 'item1: value1']) 115 116 moblab_rpc_interface.os = self.mox.CreateMockAnything() 117 moblab_rpc_interface.os.path = self.mox.CreateMockAnything() 118 moblab_rpc_interface.os.path.exists( 119 moblab_rpc_interface._CONFIG.shadow_file).AndReturn( 120 True) 121 mockShadowFile = self.mox.CreateMockAnything() 122 mockShadowFileContents = StringIO.StringIO() 123 mockShadowFile.__enter__().AndReturn(mockShadowFileContents) 124 mockShadowFile.__exit__(mox.IgnoreArg(), mox.IgnoreArg(), 125 mox.IgnoreArg()) 126 open(moblab_rpc_interface._CONFIG.shadow_file, 127 'w').AndReturn(mockShadowFile) 128 moblab_rpc_interface.os.system('sudo reboot') 129 130 self.mox.ReplayAll() 131 moblab_rpc_interface.update_config_handler( 132 {'section1' : [('item1', 'value1'), 133 ('item2', 'value2')], 134 'section2' : [('item3', 'value3'), 135 ('item4', 'value4')]}) 136 137 # item1 should not be in the new shadow config as its updated value 138 # matches the original config's value. 139 self.assertEquals( 140 mockShadowFileContents.getvalue(), 141 '[section2]\nitem3 = value3\nitem4 = value4\n\n' 142 '[section1]\nitem2 = value2\n\n') 143 144 145 def testResetConfig(self): 146 """Ensure that reset opens the shadow_config file for writing.""" 147 self.setIsMoblab(True) 148 config_mock = self.mox.CreateMockAnything() 149 moblab_rpc_interface._CONFIG = config_mock 150 config_mock.shadow_file = 'shadow_config.ini' 151 self.mox.StubOutWithMock(__builtin__, 'open') 152 mockFile = self.mox.CreateMockAnything() 153 file_contents = self.mox.CreateMockAnything() 154 mockFile.__enter__().AndReturn(file_contents) 155 mockFile.__exit__(mox.IgnoreArg(), mox.IgnoreArg(), mox.IgnoreArg()) 156 open(config_mock.shadow_file, 'w').AndReturn(mockFile) 157 moblab_rpc_interface.os = self.mox.CreateMockAnything() 158 moblab_rpc_interface.os.system('sudo reboot') 159 self.mox.ReplayAll() 160 moblab_rpc_interface.reset_config_settings() 161 162 163 def testSetBotoKey(self): 164 """Ensure that the botokey path supplied is copied correctly.""" 165 self.setIsMoblab(True) 166 boto_key = '/tmp/boto' 167 moblab_rpc_interface.os.path = self.mox.CreateMockAnything() 168 moblab_rpc_interface.os.path.exists(boto_key).AndReturn( 169 True) 170 moblab_rpc_interface.shutil = self.mox.CreateMockAnything() 171 moblab_rpc_interface.shutil.copyfile( 172 boto_key, moblab_rpc_interface.MOBLAB_BOTO_LOCATION) 173 self.mox.ReplayAll() 174 moblab_rpc_interface.set_boto_key(boto_key) 175 176 177 def testSetLaunchControlKey(self): 178 """Ensure that the Launch Control key path supplied is copied correctly. 179 """ 180 self.setIsMoblab(True) 181 launch_control_key = '/tmp/launch_control' 182 moblab_rpc_interface.os = self.mox.CreateMockAnything() 183 moblab_rpc_interface.os.path = self.mox.CreateMockAnything() 184 moblab_rpc_interface.os.path.exists(launch_control_key).AndReturn( 185 True) 186 moblab_rpc_interface.shutil = self.mox.CreateMockAnything() 187 moblab_rpc_interface.shutil.copyfile( 188 launch_control_key, 189 moblab_host.MOBLAB_LAUNCH_CONTROL_KEY_LOCATION) 190 moblab_rpc_interface.os.system('sudo restart moblab-devserver-init') 191 self.mox.ReplayAll() 192 moblab_rpc_interface.set_launch_control_key(launch_control_key) 193 194 195 def testGetNetworkInfo(self): 196 """Ensure the network info is properly converted to a dict.""" 197 self.setIsMoblab(True) 198 199 self.mox.StubOutWithMock(moblab_rpc_interface, '_get_network_info') 200 moblab_rpc_interface._get_network_info().AndReturn(('10.0.0.1', True)) 201 self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization') 202 203 rpc_utils.prepare_for_serialization( 204 {'is_connected': True, 'server_ips': ['10.0.0.1']}) 205 self.mox.ReplayAll() 206 moblab_rpc_interface.get_network_info() 207 self.mox.VerifyAll() 208 209 210 def testGetNetworkInfoWithNoIp(self): 211 """Queries network info with no public IP address.""" 212 self.setIsMoblab(True) 213 214 self.mox.StubOutWithMock(moblab_rpc_interface, '_get_network_info') 215 moblab_rpc_interface._get_network_info().AndReturn((None, False)) 216 self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization') 217 218 rpc_utils.prepare_for_serialization( 219 {'is_connected': False}) 220 self.mox.ReplayAll() 221 moblab_rpc_interface.get_network_info() 222 self.mox.VerifyAll() 223 224 225 def testGetNetworkInfoWithNoConnectivity(self): 226 """Queries network info with public IP address but no connectivity.""" 227 self.setIsMoblab(True) 228 229 self.mox.StubOutWithMock(moblab_rpc_interface, '_get_network_info') 230 moblab_rpc_interface._get_network_info().AndReturn(('10.0.0.1', False)) 231 self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization') 232 233 rpc_utils.prepare_for_serialization( 234 {'is_connected': False, 'server_ips': ['10.0.0.1']}) 235 self.mox.ReplayAll() 236 moblab_rpc_interface.get_network_info() 237 self.mox.VerifyAll() 238 239 240 def testGetCloudStorageInfo(self): 241 """Ensure the cloud storage info is properly converted to a dict.""" 242 self.setIsMoblab(True) 243 config_mock = self.mox.CreateMockAnything() 244 moblab_rpc_interface._CONFIG = config_mock 245 config_mock.get_config_value( 246 'CROS', 'image_storage_server').AndReturn('gs://bucket1') 247 config_mock.get_config_value( 248 'CROS', 'results_storage_server', default=None).AndReturn( 249 'gs://bucket2') 250 self.mox.StubOutWithMock(moblab_rpc_interface, '_get_boto_config') 251 moblab_rpc_interface._get_boto_config().AndReturn(config_mock) 252 config_mock.sections().AndReturn(['Credentials', 'b']) 253 config_mock.options('Credentials').AndReturn( 254 ['gs_access_key_id', 'gs_secret_access_key']) 255 config_mock.get( 256 'Credentials', 'gs_access_key_id').AndReturn('key') 257 config_mock.get( 258 'Credentials', 'gs_secret_access_key').AndReturn('secret') 259 rpc_utils.prepare_for_serialization( 260 { 261 'gs_access_key_id': 'key', 262 'gs_secret_access_key' : 'secret', 263 'use_existing_boto_file': True, 264 'image_storage_server' : 'gs://bucket1', 265 'results_storage_server' : 'gs://bucket2' 266 }) 267 self.mox.ReplayAll() 268 moblab_rpc_interface.get_cloud_storage_info() 269 self.mox.VerifyAll() 270 271 272 def testValidateCloudStorageInfo(self): 273 """ Ensure the cloud storage info validation flow.""" 274 self.setIsMoblab(True) 275 cloud_storage_info = { 276 'use_existing_boto_file': False, 277 'gs_access_key_id': 'key', 278 'gs_secret_access_key': 'secret', 279 'image_storage_server': 'gs://bucket1', 280 'results_storage_server': 'gs://bucket2'} 281 self.mox.StubOutWithMock(moblab_rpc_interface, '_is_valid_boto_key') 282 self.mox.StubOutWithMock(moblab_rpc_interface, '_is_valid_bucket') 283 moblab_rpc_interface._is_valid_boto_key( 284 'key', 'secret').AndReturn((True, None)) 285 moblab_rpc_interface._is_valid_bucket( 286 'key', 'secret', 'bucket1').AndReturn((True, None)) 287 moblab_rpc_interface._is_valid_bucket( 288 'key', 'secret', 'bucket2').AndReturn((True, None)) 289 rpc_utils.prepare_for_serialization( 290 {'status_ok': True }) 291 self.mox.ReplayAll() 292 moblab_rpc_interface.validate_cloud_storage_info(cloud_storage_info) 293 self.mox.VerifyAll() 294 295 296 def testGetBucketNameFromUrl(self): 297 """Gets bucket name from bucket URL.""" 298 self.assertEquals( 299 'bucket_name-123', 300 moblab_rpc_interface._get_bucket_name_from_url( 301 'gs://bucket_name-123')) 302 self.assertEquals( 303 'bucket_name-123', 304 moblab_rpc_interface._get_bucket_name_from_url( 305 'gs://bucket_name-123/')) 306 self.assertEquals( 307 'bucket_name-123', 308 moblab_rpc_interface._get_bucket_name_from_url( 309 'gs://bucket_name-123/a/b/c')) 310 self.assertIsNone(moblab_rpc_interface._get_bucket_name_from_url( 311 'bucket_name-123/a/b/c')) 312 313 314 def testIsValidBotoKeyValid(self): 315 """Tests the boto key validation flow.""" 316 if boto is None: 317 logging.info('skip test since boto module not installed') 318 return 319 conn = self.mox.CreateMockAnything() 320 self.mox.StubOutWithMock(boto, 'connect_gs') 321 boto.connect_gs('key', 'secret').AndReturn(conn) 322 conn.get_all_buckets().AndReturn(['a', 'b']) 323 conn.close() 324 self.mox.ReplayAll() 325 valid, details = moblab_rpc_interface._is_valid_boto_key('key', 'secret') 326 self.assertTrue(valid) 327 self.mox.VerifyAll() 328 329 330 def testIsValidBotoKeyInvalid(self): 331 """Tests the boto key validation with invalid key.""" 332 if boto is None: 333 logging.info('skip test since boto module not installed') 334 return 335 conn = self.mox.CreateMockAnything() 336 self.mox.StubOutWithMock(boto, 'connect_gs') 337 boto.connect_gs('key', 'secret').AndReturn(conn) 338 conn.get_all_buckets().AndRaise( 339 boto.exception.GSResponseError('bad', 'reason')) 340 conn.close() 341 self.mox.ReplayAll() 342 valid, details = moblab_rpc_interface._is_valid_boto_key('key', 'secret') 343 self.assertFalse(valid) 344 self.assertEquals('The boto access key is not valid', details) 345 self.mox.VerifyAll() 346 347 348 def testIsValidBucketValid(self): 349 """Tests the bucket vaildation flow.""" 350 if boto is None: 351 logging.info('skip test since boto module not installed') 352 return 353 conn = self.mox.CreateMockAnything() 354 self.mox.StubOutWithMock(boto, 'connect_gs') 355 boto.connect_gs('key', 'secret').AndReturn(conn) 356 conn.lookup('bucket').AndReturn('bucket') 357 conn.close() 358 self.mox.ReplayAll() 359 valid, details = moblab_rpc_interface._is_valid_bucket( 360 'key', 'secret', 'bucket') 361 self.assertTrue(valid) 362 self.mox.VerifyAll() 363 364 365 def testIsValidBucketInvalid(self): 366 """Tests the bucket validation flow with invalid key.""" 367 if boto is None: 368 logging.info('skip test since boto module not installed') 369 return 370 conn = self.mox.CreateMockAnything() 371 self.mox.StubOutWithMock(boto, 'connect_gs') 372 boto.connect_gs('key', 'secret').AndReturn(conn) 373 conn.lookup('bucket').AndReturn(None) 374 conn.close() 375 self.mox.ReplayAll() 376 valid, details = moblab_rpc_interface._is_valid_bucket( 377 'key', 'secret', 'bucket') 378 self.assertFalse(valid) 379 self.assertEquals("Bucket bucket does not exist.", details) 380 self.mox.VerifyAll() 381 382 383 def testGetShadowConfigFromPartialUpdate(self): 384 """Tests getting shadow configuration based on partial upate.""" 385 partial_config = { 386 'section1': [ 387 ('opt1', 'value1'), 388 ('opt2', 'value2'), 389 ('opt3', 'value3'), 390 ('opt4', 'value4'), 391 ] 392 } 393 shadow_config_str = "[section1]\nopt2 = value2_1\nopt4 = value4_1" 394 shadow_config = ConfigParser.ConfigParser() 395 shadow_config.readfp(StringIO.StringIO(shadow_config_str)) 396 original_config = self.mox.CreateMockAnything() 397 self.mox.StubOutWithMock(moblab_rpc_interface, '_read_original_config') 398 self.mox.StubOutWithMock(moblab_rpc_interface, '_read_raw_config') 399 moblab_rpc_interface._read_original_config().AndReturn(original_config) 400 moblab_rpc_interface._read_raw_config( 401 moblab_rpc_interface._CONFIG.shadow_file).AndReturn(shadow_config) 402 original_config.get_config_value( 403 'section1', 'opt1', 404 allow_blank=True, default='').AndReturn('value1') 405 original_config.get_config_value( 406 'section1', 'opt2', 407 allow_blank=True, default='').AndReturn('value2') 408 original_config.get_config_value( 409 'section1', 'opt3', 410 allow_blank=True, default='').AndReturn('blah') 411 original_config.get_config_value( 412 'section1', 'opt4', 413 allow_blank=True, default='').AndReturn('blah') 414 self.mox.ReplayAll() 415 shadow_config = moblab_rpc_interface._get_shadow_config_from_partial_update( 416 partial_config) 417 # opt1 same as the original. 418 self.assertFalse(shadow_config.has_option('section1', 'opt1')) 419 # opt2 reverts back to original 420 self.assertFalse(shadow_config.has_option('section1', 'opt2')) 421 # opt3 is updated from original. 422 self.assertEquals('value3', shadow_config.get('section1', 'opt3')) 423 # opt3 in shadow but updated again. 424 self.assertEquals('value4', shadow_config.get('section1', 'opt4')) 425 self.mox.VerifyAll() 426 427 428 def testGetShadowConfigFromPartialUpdateWithNewSection(self): 429 """ 430 Test getting shadown configuration based on partial update with new section. 431 """ 432 partial_config = { 433 'section2': [ 434 ('opt5', 'value5'), 435 ('opt6', 'value6'), 436 ], 437 } 438 shadow_config_str = "[section1]\nopt2 = value2_1\n" 439 shadow_config = ConfigParser.ConfigParser() 440 shadow_config.readfp(StringIO.StringIO(shadow_config_str)) 441 original_config = self.mox.CreateMockAnything() 442 self.mox.StubOutWithMock(moblab_rpc_interface, '_read_original_config') 443 self.mox.StubOutWithMock(moblab_rpc_interface, '_read_raw_config') 444 moblab_rpc_interface._read_original_config().AndReturn(original_config) 445 moblab_rpc_interface._read_raw_config( 446 moblab_rpc_interface._CONFIG.shadow_file).AndReturn(shadow_config) 447 original_config.get_config_value( 448 'section2', 'opt5', 449 allow_blank=True, default='').AndReturn('value5') 450 original_config.get_config_value( 451 'section2', 'opt6', 452 allow_blank=True, default='').AndReturn('blah') 453 self.mox.ReplayAll() 454 shadow_config = moblab_rpc_interface._get_shadow_config_from_partial_update( 455 partial_config) 456 # opt2 is still in shadow 457 self.assertEquals('value2_1', shadow_config.get('section1', 'opt2')) 458 # opt5 is not changed. 459 self.assertFalse(shadow_config.has_option('section2', 'opt5')) 460 # opt6 is updated. 461 self.assertEquals('value6', shadow_config.get('section2', 'opt6')) 462 self.mox.VerifyAll() 463 464 465 if __name__ == '__main__': 466 unittest.main() 467
