aosp12/external/autotest/site_utils/pubsub_utils_unittest.py

#!/usr/bin/env python2
# Copyright 2016 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.

"""Unit test for pubsub_utils.py"""

from __future__ import print_function
import os
import unittest

import mox

# TODO(crbug.com/1050892): The unittests rely on apiclient in chromite.
import chromite  # pylint: disable=unused-import

from apiclient import discovery
from oauth2client.client import ApplicationDefaultCredentialsError
from oauth2client.client import GoogleCredentials
from googleapiclient.errors import UnknownApiNameOrVersion

import pubsub_utils

_TEST_CLOUD_SERVICE_ACCOUNT_FILE = '/tmp/test-credential'


class MockedPubSub(object):
    """A mocked PubSub handle."""
    def __init__(self, test, topic, msg, retry, ret_val=None,
                 raise_except=False):
        self.test = test
        self.topic = topic
        self.msg = msg
        self.retry = retry
        self.ret_val = ret_val
        self.raise_except = raise_except

    def projects(self):
        """Mocked PubSub projects."""
        return self

    def topics(self):
        """Mocked PubSub topics."""
        return self

    def publish(self, topic, body):
        """Mocked PubSub publish method.

        @param topic: PubSub topic string.
        @param body: PubSub notification body.
        """
        self.test.assertEquals(self.topic, topic)
        self.test.assertEquals(self.msg, body['messages'][0])
        return self

    def execute(self, num_retries):
        """Mocked PubSub execute method.

        @param num_retries: Number of retries.
        """
        self.test.assertEquals(self.retry, num_retries)
        if self.raise_except:
            raise Exception()
        return self.ret_val


def _create_sample_message():
    """Creates a sample pubsub message."""
    msg_payload = {'data': 'sample data'}
    msg_attributes = {}
    msg_attributes['var'] = 'value'
    msg_payload['attributes'] = msg_attributes

    return msg_payload


class PubSubTests(mox.MoxTestBase):
    """Tests for pubsub related functios."""

    def test_pubsub_with_no_service_account(self):
        """Test getting the pubsub service"""
        self.mox.StubOutWithMock(os.path, 'isfile')
        self.mox.ReplayAll()
        with self.assertRaises(pubsub_utils.PubSubException):
            pubsub_utils.PubSubClient()
        self.mox.VerifyAll()

    def test_pubsub_with_non_existing_service_account(self):
        """Test getting the pubsub service"""
        self.mox.StubOutWithMock(os.path, 'isfile')
        os.path.isfile(_TEST_CLOUD_SERVICE_ACCOUNT_FILE).AndReturn(False)
        self.mox.ReplayAll()
        with self.assertRaises(pubsub_utils.PubSubException):
            pubsub_utils.PubSubClient(_TEST_CLOUD_SERVICE_ACCOUNT_FILE)
        self.mox.VerifyAll()

    def test_pubsub_with_corrupted_service_account(self):
        """Test pubsub with corrupted service account."""
        self.mox.StubOutWithMock(os.path, 'isfile')
        self.mox.StubOutWithMock(GoogleCredentials, 'from_stream')
        os.path.isfile(_TEST_CLOUD_SERVICE_ACCOUNT_FILE).AndReturn(True)
        GoogleCredentials.from_stream(
            _TEST_CLOUD_SERVICE_ACCOUNT_FILE).AndRaise(
                ApplicationDefaultCredentialsError())
        self.mox.ReplayAll()
        with self.assertRaises(pubsub_utils.PubSubException):
            pubsub_utils.PubSubClient(_TEST_CLOUD_SERVICE_ACCOUNT_FILE)
        self.mox.VerifyAll()

    def test_pubsub_with_invalid_service_account(self):
        """Test pubsubwith invalid service account."""
        self.mox.StubOutWithMock(os.path, 'isfile')
        self.mox.StubOutWithMock(GoogleCredentials, 'from_stream')
        os.path.isfile(_TEST_CLOUD_SERVICE_ACCOUNT_FILE).AndReturn(True)
        credentials = self.mox.CreateMock(GoogleCredentials)
        GoogleCredentials.from_stream(
            _TEST_CLOUD_SERVICE_ACCOUNT_FILE).AndReturn(credentials)
        credentials.create_scoped_required().AndReturn(True)
        credentials.create_scoped(pubsub_utils.PUBSUB_SCOPES).AndReturn(
            credentials)
        self.mox.StubOutWithMock(discovery, 'build')
        discovery.build(
            pubsub_utils.PUBSUB_SERVICE_NAME,
            pubsub_utils.PUBSUB_VERSION,
            credentials=credentials).AndRaise(UnknownApiNameOrVersion())
        self.mox.ReplayAll()
        with self.assertRaises(pubsub_utils.PubSubException):
            msg = _create_sample_message()
            pubsub_client = pubsub_utils.PubSubClient(
                _TEST_CLOUD_SERVICE_ACCOUNT_FILE)
            pubsub_client.publish_notifications('test_topic', [msg])
        self.mox.VerifyAll()

    def test_publish_notifications(self):
        """Test getting the pubsub service"""
        self.mox.StubOutWithMock(os.path, 'isfile')
        self.mox.StubOutWithMock(GoogleCredentials, 'from_stream')
        os.path.isfile(_TEST_CLOUD_SERVICE_ACCOUNT_FILE).AndReturn(True)
        credentials = self.mox.CreateMock(GoogleCredentials)
        GoogleCredentials.from_stream(
            _TEST_CLOUD_SERVICE_ACCOUNT_FILE).AndReturn(credentials)
        credentials.create_scoped_required().AndReturn(True)
        credentials.create_scoped(pubsub_utils.PUBSUB_SCOPES).AndReturn(
            credentials)
        self.mox.StubOutWithMock(discovery, 'build')
        msg = _create_sample_message()
        discovery.build(
            pubsub_utils.PUBSUB_SERVICE_NAME,
            pubsub_utils.PUBSUB_VERSION,
            credentials=credentials).AndReturn(MockedPubSub(
                self,
                'test_topic',
                msg,
                pubsub_utils.DEFAULT_PUBSUB_NUM_RETRIES,
                # use tuple ('123') instead of list just for easy to
                # write the test.
                ret_val={'messageIds': ('123')}))

        self.mox.ReplayAll()
        pubsub_client = pubsub_utils.PubSubClient(
                _TEST_CLOUD_SERVICE_ACCOUNT_FILE)
        msg_ids = pubsub_client.publish_notifications('test_topic', [msg])
        self.assertEquals(('123'), msg_ids)

        self.mox.VerifyAll()


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