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.
13try:
14    import boto
15except ImportError:
16    boto = None
17import ConfigParser
18import logging
19import mox
20import StringIO
21import unittest
22
23import common
24
25from autotest_lib.client.common_lib import error
26from autotest_lib.client.common_lib import global_config
27from autotest_lib.client.common_lib import lsbrelease_utils
28from autotest_lib.frontend import setup_django_environment
29from autotest_lib.frontend.afe import frontend_test_utils
30from autotest_lib.frontend.afe import moblab_rpc_interface
31from autotest_lib.frontend.afe import rpc_utils
32from autotest_lib.server import utils
33from autotest_lib.server.hosts import moblab_host
34
35
36class 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
465if __name__ == '__main__':
466    unittest.main()
467