blob: c052a9f9d296afa11131250dc02535c959f673b2 [file] [log] [blame] [edit]
#!/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()