xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 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. |
xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 13 | import ConfigParser |
xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 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 |
Keith Haddow | 350d789 | 2017-04-27 15:47:53 -0700 | [diff] [blame] | 29 | from autotest_lib.client.common_lib import utils as common_lib_utils |
xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 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 | |
xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 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'} |
Keith Haddow | 350d789 | 2017-04-27 15:47:53 -0700 | [diff] [blame] | 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 }) |
xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 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 | |
xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 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 | |
Keith Haddow | 3102fd4 | 2017-05-10 11:55:29 -0700 | [diff] [blame] | 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 | |
Keith Haddow | c5ec060 | 2017-06-01 16:49:17 -0700 | [diff] [blame] | 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', |
Keith Haddow | 3102fd4 | 2017-05-10 11:55:29 -0700 | [diff] [blame] | 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() |
xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 425 | |
Keith Haddow | 350d789 | 2017-04-27 15:47:53 -0700 | [diff] [blame] | 426 | def testRunBucketPerformanceTestFail(self): |
Keith Haddow | c5ec060 | 2017-06-01 16:49:17 -0700 | [diff] [blame] | 427 | self.mox.StubOutWithMock(moblab_rpc_interface.GsUtil, 'get_gsutil_cmd') |
| 428 | moblab_rpc_interface.GsUtil.get_gsutil_cmd().AndReturn( |
| 429 | '/path/to/gsutil') |
Keith Haddow | 350d789 | 2017-04-27 15:47:53 -0700 | [diff] [blame] | 430 | self.mox.StubOutWithMock(common_lib_utils, 'run') |
Keith Haddow | c5ec060 | 2017-06-01 16:49:17 -0700 | [diff] [blame] | 431 | common_lib_utils.run('/path/to/gsutil', |
Keith Haddow | 350d789 | 2017-04-27 15:47:53 -0700 | [diff] [blame] | 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 | |
Keith Haddow | 21e33a5 | 2017-05-17 15:42:58 -0700 | [diff] [blame] | 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/') |
Keith Haddow | c5ec060 | 2017-06-01 16:49:17 -0700 | [diff] [blame] | 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 | |
Keith Haddow | 21e33a5 | 2017-05-17 15:42:58 -0700 | [diff] [blame] | 460 | self.mox.StubOutWithMock(common_lib_utils, 'run') |
Keith Haddow | c5ec060 | 2017-06-01 16:49:17 -0700 | [diff] [blame] | 461 | common_lib_utils.run('/path/to/gsutil', |
Keith Haddow | 21e33a5 | 2017-05-17 15:42:58 -0700 | [diff] [blame] | 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/') |
Keith Haddow | c5ec060 | 2017-06-01 16:49:17 -0700 | [diff] [blame] | 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 | |
Keith Haddow | 21e33a5 | 2017-05-17 15:42:58 -0700 | [diff] [blame] | 478 | self.mox.StubOutWithMock(common_lib_utils, 'run') |
Keith Haddow | c5ec060 | 2017-06-01 16:49:17 -0700 | [diff] [blame] | 479 | common_lib_utils.run('/path/to/gsutil', |
Keith Haddow | 21e33a5 | 2017-05-17 15:42:58 -0700 | [diff] [blame] | 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 | |
Matt Mallett | b0f8dc7 | 2018-01-25 15:15:58 -0800 | [diff] [blame] | 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']) |
Matt Mallett | b0f8dc7 | 2018-01-25 15:15:58 -0800 | [diff] [blame] | 499 | moblab_rpc_interface.subprocess.check_call(['sudo', |
Matt Mallett | c223505 | 2018-01-30 16:47:08 -0800 | [diff] [blame] | 500 | update_engine_client, '--is_reboot_needed']) |
Matt Mallett | b0f8dc7 | 2018-01-25 15:15:58 -0800 | [diff] [blame] | 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 | |
Matt Mallett | 1674447 | 2018-02-08 10:23:03 -0800 | [diff] [blame] | 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 | |
Matt Mallett | b0f8dc7 | 2018-01-25 15:15:58 -0800 | [diff] [blame] | 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 | |
Matt Mallett | b879766 | 2018-02-01 15:34:28 -0800 | [diff] [blame] | 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 | ) |
Matt Mallett | b0f8dc7 | 2018-01-25 15:15:58 -0800 | [diff] [blame] | 572 | self.mox.ReplayAll() |
| 573 | moblab_rpc_interface._check_for_system_update() |
| 574 | |
Matt Mallett | 31c1eec | 2018-01-30 16:13:26 -0800 | [diff] [blame] | 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 | |
Matt Mallett | f5321f0 | 2018-02-23 16:40:23 -0800 | [diff] [blame] | 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) |
Keith Haddow | 350d789 | 2017-04-27 15:47:53 -0700 | [diff] [blame] | 705 | |
Matt Mallett | 3b5a706 | 2018-02-05 10:15:27 -0800 | [diff] [blame] | 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 | |
xixuan | ba232a3 | 2016-08-25 17:01:59 -0700 | [diff] [blame] | 728 | if __name__ == '__main__': |
| 729 | unittest.main() |