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
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.
13import ConfigParser
14import 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
29from autotest_lib.client.common_lib import utils as common_lib_utils
30
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
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
494if __name__ == '__main__':
495    unittest.main()
496