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