site_utils/cloud_console_client_unittest.py

#!/usr/bin/python
# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import base64
import mox
import unittest

import common

from autotest_lib.client.common_lib import utils
from autotest_lib.site_utils import cloud_console_client
from autotest_lib.site_utils import cloud_console_pb2 as cpcon
from autotest_lib.site_utils import pubsub_utils

class PubSubBasedClientTests(mox.MoxTestBase):
    """Tests for the 'PubSubBasedClient'."""

    def setUp(self):
        super(PubSubBasedClientTests, self).setUp()
        self._pubsub_client_mock = self.mox.CreateMock(
                pubsub_utils.PubSubClient)
        self._stubs = mox.stubout.StubOutForTesting()
        self._stubs.Set(cloud_console_client, '_create_pubsub_client',
                lambda x: self._pubsub_client_mock)

    def tearDown(self):
        self._stubs.UnsetAll()
        super(PubSubBasedClientTests, self).tearDown()

    def test_create_test_result_notification(self):
        """Tests the test result notification message."""
        self._console_client = cloud_console_client.PubSubBasedClient()
        self.mox.StubOutWithMock(utils,
                                 'get_moblab_serial_number')
        utils.get_moblab_serial_number().AndReturn(
            'PV120BB8JAC01E')
        self.mox.StubOutWithMock(utils, 'get_moblab_id')
        utils.get_moblab_id().AndReturn(
            'c8386d92-9ad1-11e6-80f5-111111111111')
        self.mox.ReplayAll()
        console_client = cloud_console_client.PubSubBasedClient()
        msg = console_client._create_test_job_offloaded_message(
                'gs://test_bucket/123-moblab')
        self.assertEquals(base64.b64encode(
            cloud_console_client.LEGACY_TEST_OFFLOAD_MESSAGE), msg['data'])
        self.assertEquals(
            cloud_console_client.CURRENT_MESSAGE_VERSION,
            msg['attributes'][cloud_console_client.LEGACY_ATTR_VERSION])
        self.assertEquals(
            '1c:dc:d1:11:01:e1',
            msg['attributes'][cloud_console_client.LEGACY_ATTR_MOBLAB_MAC]
            )
        self.assertEquals(
            'c8386d92-9ad1-11e6-80f5-111111111111',
            msg['attributes'][cloud_console_client.LEGACY_ATTR_MOBLAB_ID])
        self.assertEquals(
            'gs://test_bucket/123-moblab',
            msg['attributes'][cloud_console_client.LEGACY_ATTR_GCS_URI])
        self.mox.VerifyAll()

    def test_send_test_job_offloaded_message(self):
        """Tests send job offloaded notification."""
        console_client = cloud_console_client.PubSubBasedClient(
                pubsub_topic='test topic')

        message = {'data': 'dummy data', 'attributes' : {'key' : 'value'}}
        self.mox.StubOutWithMock(cloud_console_client.PubSubBasedClient,
                '_create_test_job_offloaded_message')
        console_client._create_test_job_offloaded_message(
                'gs://test_bucket/123-moblab').AndReturn(message)

        msg_ids = ['1']
        self._pubsub_client_mock.publish_notifications(
                'test topic', [message]).AndReturn(msg_ids)

        self.mox.ReplayAll()
        result = console_client.send_test_job_offloaded_message(
                'gs://test_bucket/123-moblab')
        self.assertTrue(result)
        self.mox.VerifyAll()

    def test_send_heartbeat(self):
        """Tests send heartbeat."""
        console_client = cloud_console_client.PubSubBasedClient(
                pubsub_topic='test topic')
        self.mox.StubOutWithMock(utils, 'get_moblab_id')
        utils.get_moblab_id().AndReturn(
            'c8386d92-9ad1-11e6-80f5-111111111111')

        message = {
                'attributes' : {
                    'ATTR_MOBLAB_ID': 'c8386d92-9ad1-11e6-80f5-111111111111',
                    'ATTR_MESSAGE_VERSION': '1',
                    'ATTR_MESSAGE_TYPE': 'MSG_MOBLAB_HEARTBEAT',
                    'ATTR_MOBLAB_MAC_ADDRESS': '8c:dc:d4:56:06:e7'}}
        msg_ids = ['1']
        self._pubsub_client_mock.publish_notifications(
                'test topic', [message]).AndReturn(msg_ids)
        self.mox.ReplayAll()
        console_client.send_heartbeat()
        self.mox.VerifyAll()

    def test_send_event(self):
        """Tests send heartbeat."""
        console_client = cloud_console_client.PubSubBasedClient(
                pubsub_topic='test topic')
        self.mox.StubOutWithMock(utils, 'get_moblab_id')
        utils.get_moblab_id().AndReturn(
            'c8386d92-9ad1-11e6-80f5-111111111111')

        message = {
                'data': '\x08\x01\x12\x0ethis is a test',
                'attributes' : {
                    'ATTR_MOBLAB_ID': 'c8386d92-9ad1-11e6-80f5-111111111111',
                    'ATTR_MESSAGE_VERSION': '1',
                    'ATTR_MESSAGE_TYPE': 'MSG_MOBLAB_REMOTE_EVENT',
                    'ATTR_MOBLAB_MAC_ADDRESS': '8c:dc:d4:56:06:e7'}}
        msg_ids = ['1']
        self._pubsub_client_mock.publish_notifications(
                'test topic', [message]).AndReturn(msg_ids)
        self.mox.ReplayAll()
        console_client.send_event(
                cpcon.RemoteEventMessage.EVENT_MOBLAB_BOOT_COMPLETE,
                'this is a test')
        self.mox.VerifyAll()

if __name__ == '__main__':
    unittest.main()