Wrap the pubsub into a cloud console client.
BUG=chromium:724523
TEST=unit tests.
Change-Id: I4338cce8c2c983b4cd50b71014d4f1ca5667cd78
Reviewed-on: https://chromium-review.googlesource.com/540137
Commit-Ready: Michael Tang <ntang@chromium.org>
Tested-by: Michael Tang <ntang@chromium.org>
Reviewed-by: Keith Haddow <haddowk@chromium.org>
Reviewed-by: Michael Tang <ntang@chromium.org>
diff --git a/site_utils/cloud_console_client.py b/site_utils/cloud_console_client.py
new file mode 100644
index 0000000..892f677
--- /dev/null
+++ b/site_utils/cloud_console_client.py
@@ -0,0 +1,189 @@
+# 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.
+
+"""Chrome OS Parnter Concole remote actions."""
+
+from __future__ import print_function
+
+import base64
+import logging
+
+import common
+
+from autotest_lib.client.common_lib import global_config
+from autotest_lib.client.common_lib import utils
+from autotest_lib.server.hosts import moblab_host
+from autotest_lib.site_utils import pubsub_utils
+
+
+_PUBSUB_TOPIC = global_config.global_config.get_config_value(
+ 'CROS', 'cloud_notification_topic', default=None)
+
+# Test upload pubsub notification attributes
+NOTIFICATION_ATTR_VERSION = 'version'
+NOTIFICATION_ATTR_GCS_URI = 'gcs_uri'
+NOTIFICATION_ATTR_MOBLAB_MAC = 'moblab_mac_address'
+NOTIFICATION_ATTR_MOBLAB_ID = 'moblab_id'
+NOTIFICATION_VERSION = '1'
+# the message data for new test result notification.
+NEW_TEST_RESULT_MESSAGE = 'NEW_TEST_RESULT'
+
+
+ALERT_CRITICAL = 'Critical'
+ALERT_MAJOR = 'Major'
+ALERT_MINOR = 'Minor'
+
+LOG_INFO = 'Info'
+LOG_WARNING = 'Warning'
+LOG_SEVERE = 'Severe'
+LOG_FATAL = 'Fatal'
+
+
+def is_cloud_notification_enabled():
+ """Checks if cloud pubsub notification is enabled.
+
+ @returns: True if cloud pubsub notification is enabled. Otherwise, False.
+ """
+ return global_config.global_config.get_config_value(
+ 'CROS', 'cloud_notification_enabled', type=bool, default=False)
+
+
+class CloudConsoleClient(object):
+ """The remote interface to the Cloud Console."""
+ def send_heartbeat(self):
+ """Sends a heartbeat.
+
+ @returns True if the notification is successfully sent.
+ Otherwise, False.
+ """
+ pass
+
+ def send_event(self, event_type=None, event_data=None):
+ """Sends an event notification to the remote console.
+
+ @param event_type: The event type that is defined in the protobuffer
+ file 'cloud_console.proto'.
+ @param event_data: The event data.
+
+ @returns True if the notification is successfully sent.
+ Otherwise, False.
+ """
+ pass
+
+ def send_log(self, msg, level=LOG_INFO, session_id=None):
+ """Sends a log message to the remote console.
+
+ @param msg: The log message.
+ @param level: The logging level as string.
+ @param session_id: The current session id.
+
+ @returns True if the notification is successfully sent.
+ Otherwise, False.
+ """
+ pass
+
+ def send_alert(self, msg, level=ALERT_MINOR, session_id=None):
+ """Sends an alert to the remote console.
+
+ @param msg: The alert message.
+ @param level: The logging level as string.
+ @param session_id: The current session id.
+
+ @returns True if the notification is successfully sent.
+ Otherwise, False.
+ """
+ pass
+
+ def send_test_job_offloaded_message(self, gcs_uri):
+ """Sends a test job offloaded message to the remote console.
+
+ @param gcs_uri: The test result Google Cloud Storage URI.
+
+ @returns True if the notification is successfully sent.
+ Otherwise, False.
+ """
+ pass
+
+
+# Make it easy to mock out
+def _create_pubsub_client(credential):
+ return pubsub_utils.PubSubClient(credential)
+
+
+class PubSubBasedClient(CloudConsoleClient):
+ """A Cloud PubSub based implementation of the CloudConsoleClient interface.
+ """
+ def __init__(
+ self,
+ credential=moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION,
+ pubsub_topic=_PUBSUB_TOPIC):
+ """Constructor.
+
+ @param credential: The service account credential filename. Default to
+ '/home/moblab/.service_account.json'.
+ @param pubsub_topic: The cloud pubsub topic name to use.
+ """
+ super(PubSubBasedClient, self).__init__()
+ self._pubsub_client = _create_pubsub_client(credential)
+ self._pubsub_topic = pubsub_topic
+
+
+ def _create_notification_message(self, data, msg_attributes):
+ """Creates a cloud pubsub notification object.
+
+ @param data: The message data as a string.
+ @param msg_attributes: The message attribute map.
+
+ @returns: A pubsub message object with data and attributes.
+ """
+ message = {'data': data}
+ message['attributes'] = msg_attributes
+ return message
+
+ def _create_notification_attributes(self):
+ """Creates a cloud pubsub notification message attribute map.
+
+ Fills in the version, moblab mac address, and moblab id information
+ as attributes.
+
+ @returns: A pubsub messsage attribute map.
+ """
+ msg_attributes = {}
+ msg_attributes[NOTIFICATION_ATTR_VERSION] = NOTIFICATION_VERSION
+ msg_attributes[NOTIFICATION_ATTR_MOBLAB_MAC] = (
+ utils.get_default_interface_mac_address())
+ msg_attributes[NOTIFICATION_ATTR_MOBLAB_ID] = utils.get_moblab_id()
+ return msg_attributes
+
+ def _create_test_result_notification(self, gcs_uri):
+ """Construct a test result notification.
+
+ @param gcs_uri: The test result Google Cloud Storage URI.
+
+ @returns The notification message.
+ """
+ data = base64.b64encode(NEW_TEST_RESULT_MESSAGE)
+ msg_attributes = self._create_notification_attributes()
+ msg_attributes[NOTIFICATION_ATTR_GCS_URI] = gcs_uri
+
+ return self._create_notification_message(data, msg_attributes)
+
+
+ def send_test_job_offloaded_message(self, gcs_uri):
+ """Notify the cloud console a test job is offloaded.
+
+ @param gcs_uri: The test result Google Cloud Storage URI.
+
+ @returns True if the notification is successfully sent.
+ Otherwise, False.
+ """
+ logging.info('Notification on gcs_uri %s', gcs_uri)
+ message = self._create_test_result_notification(gcs_uri)
+ msg_ids = self._pubsub_client.publish_notifications(
+ self._pubsub_topic, [message])
+ if msg_ids:
+ return True
+ logging.warning('Failed to send notification on gcs_uri %s', gcs_uri)
+ return False
+