Home | History | Annotate | Download | only in site_utils
      1 #!/usr/bin/python
      2 # Copyright 2017 The Chromium OS Authors. All rights reserved.
      3 # Use of this source code is governed by a BSD-style license that can be
      4 # found in the LICENSE file.
      5 
      6 import base64
      7 import mox
      8 import unittest
      9 
     10 import common
     11 
     12 from autotest_lib.client.common_lib import utils
     13 from autotest_lib.site_utils import cloud_console_client
     14 from autotest_lib.site_utils import pubsub_utils
     15 
     16 class PubSubBasedClientTests(mox.MoxTestBase):
     17     """Tests for the 'PubSubBasedClient'."""
     18 
     19     def setUp(self):
     20         super(PubSubBasedClientTests, self).setUp()
     21         self._pubsub_client_mock = self.mox.CreateMock(
     22                 pubsub_utils.PubSubClient)
     23         self._stubs = mox.stubout.StubOutForTesting()
     24         self._stubs.Set(cloud_console_client, '_create_pubsub_client',
     25                 lambda x: self._pubsub_client_mock)
     26 
     27     def tearDown(self):
     28         self._stubs.UnsetAll()
     29         super(PubSubBasedClientTests, self).tearDown()
     30 
     31     def test_create_test_result_notification(self):
     32         """Tests the test result notification message."""
     33         self._console_client = cloud_console_client.PubSubBasedClient()
     34         self.mox.StubOutWithMock(utils, 'get_moblab_id')
     35         self.mox.StubOutWithMock(utils,
     36                                  'get_default_interface_mac_address')
     37         utils.get_default_interface_mac_address().AndReturn(
     38             '1c:dc:d1:11:01:e1')
     39         utils.get_moblab_id().AndReturn(
     40             'c8386d92-9ad1-11e6-80f5-111111111111')
     41         self.mox.ReplayAll()
     42         console_client = cloud_console_client.PubSubBasedClient()
     43         msg = console_client._create_test_result_notification(
     44                 'gs://test_bucket/123-moblab')
     45         self.assertEquals(base64.b64encode(
     46             cloud_console_client.NEW_TEST_RESULT_MESSAGE), msg['data'])
     47         self.assertEquals(
     48             cloud_console_client.NOTIFICATION_VERSION,
     49             msg['attributes'][cloud_console_client.NOTIFICATION_ATTR_VERSION])
     50         self.assertEquals(
     51             '1c:dc:d1:11:01:e1',
     52             msg['attributes'][cloud_console_client.NOTIFICATION_ATTR_MOBLAB_MAC]
     53             )
     54         self.assertEquals(
     55             'c8386d92-9ad1-11e6-80f5-111111111111',
     56             msg['attributes'][cloud_console_client.NOTIFICATION_ATTR_MOBLAB_ID])
     57         self.assertEquals(
     58             'gs://test_bucket/123-moblab',
     59             msg['attributes'][cloud_console_client.NOTIFICATION_ATTR_GCS_URI])
     60         self.mox.VerifyAll()
     61 
     62 
     63     def test_send_test_job_offloaded_message(self):
     64         """Tests send job offloaded notification."""
     65         console_client = cloud_console_client.PubSubBasedClient(
     66                 pubsub_topic='test topic')
     67 
     68         # self.mox.ResetAll()
     69         message = {'data': 'dummy data', 'attributes' : {'key' : 'value'}}
     70         self.mox.StubOutWithMock(cloud_console_client.PubSubBasedClient,
     71                 '_create_test_result_notification')
     72         console_client._create_test_result_notification(
     73                 'gs://test_bucket/123-moblab').AndReturn(message)
     74 
     75         msg_ids = ['1']
     76         self._pubsub_client_mock.publish_notifications(
     77                 'test topic', [message]).AndReturn(msg_ids)
     78 
     79         self.mox.ReplayAll()
     80         console_client.send_test_job_offloaded_message(
     81                 'gs://test_bucket/123-moblab')
     82         self.mox.VerifyAll()
     83 
     84 
     85 if __name__ == '__main__':
     86     unittest.main()
     87