# 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
|
from autotest_lib.site_utils import cloud_console_pb2 as cpcon
|
|
|
_PUBSUB_TOPIC = global_config.global_config.get_config_value(
|
'CROS', 'cloud_notification_topic', default=None)
|
|
# Current notification version.
|
CURRENT_MESSAGE_VERSION = '1'
|
|
# Test upload pubsub notification attributes
|
LEGACY_ATTR_VERSION = 'version'
|
LEGACY_ATTR_GCS_URI = 'gcs_uri'
|
LEGACY_ATTR_MOBLAB_MAC = 'moblab_mac_address'
|
LEGACY_ATTR_MOBLAB_ID = 'moblab_id'
|
# the message data for new test result notification.
|
LEGACY_TEST_OFFLOAD_MESSAGE = 'NEW_TEST_RESULT'
|
|
|
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)
|
|
|
def _get_message_type_name(message_type_enum):
|
"""Gets the message type name from message type enum.
|
|
@param message_type_enum: The message type enum.
|
|
@return The corresponding message type name as string, or 'MSG_UNKNOWN'.
|
"""
|
return cpcon.MessageType.Name(message_type_enum)
|
|
|
def _get_attribute_name(attribute_enum):
|
"""Gets the message attribute name from attribte enum.
|
|
@param attribute_enum: The attribute enum.
|
|
@return The corresponding attribute name as string, or 'ATTR_INVALID'.
|
"""
|
return cpcon.MessageAttribute.Name(attribute_enum)
|
|
|
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=None, session_id=None):
|
"""Sends a log message to the remote console.
|
|
@param msg: The log message.
|
@param level: The logging level.
|
@param session_id: The current session id.
|
|
@returns True if the notification is successfully sent.
|
Otherwise, False.
|
"""
|
pass
|
|
def send_alert(self, msg, level=None, session_id=None):
|
"""Sends an alert to the remote console.
|
|
@param msg: The alert message.
|
@param level: The logging level.
|
@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_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 = {}
|
if data:
|
message['data'] = data
|
if msg_attributes:
|
message['attributes'] = msg_attributes
|
return message
|
|
def _create_message_attributes(self, message_type_enum):
|
"""Creates a cloud pubsub notification message attribute map.
|
|
Fills in the version, moblab mac address, and moblab id information
|
as attributes.
|
|
@param message_type_enum The message type enum.
|
|
@returns: A pubsub messsage attribute map.
|
"""
|
msg_attributes = {}
|
msg_attributes[_get_attribute_name(cpcon.ATTR_MESSAGE_TYPE)] = (
|
_get_message_type_name(message_type_enum))
|
msg_attributes[_get_attribute_name(cpcon.ATTR_MESSAGE_VERSION)] = (
|
CURRENT_MESSAGE_VERSION)
|
msg_attributes[_get_attribute_name(cpcon.ATTR_MOBLAB_MAC_ADDRESS)] = (
|
utils.get_moblab_serial_number())
|
msg_attributes[_get_attribute_name(cpcon.ATTR_MOBLAB_ID)] = (
|
utils.get_moblab_id())
|
return msg_attributes
|
|
def _create_test_job_offloaded_message(self, gcs_uri):
|
"""Construct a test result notification.
|
|
TODO(ntang): switch LEGACY to new message format.
|
@param gcs_uri: The test result Google Cloud Storage URI.
|
|
@returns The notification message.
|
"""
|
data = base64.b64encode(LEGACY_TEST_OFFLOAD_MESSAGE)
|
msg_attributes = {}
|
msg_attributes[LEGACY_ATTR_VERSION] = CURRENT_MESSAGE_VERSION
|
msg_attributes[LEGACY_ATTR_MOBLAB_MAC] = (
|
utils.get_moblab_serial_number())
|
msg_attributes[LEGACY_ATTR_MOBLAB_ID] = utils.get_moblab_id()
|
msg_attributes[LEGACY_ATTR_GCS_URI] = gcs_uri
|
|
return self._create_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_job_offloaded_message(gcs_uri)
|
return self._publish_notification(message)
|
|
|
def _publish_notification(self, message):
|
msg_ids = self._pubsub_client.publish_notifications(
|
self._pubsub_topic, [message])
|
|
if msg_ids:
|
logging.debug('Successfully sent out a notification')
|
return True
|
logging.warning('Failed to send notification %s', str(message))
|
return False
|
|
def send_heartbeat(self):
|
"""Sends a heartbeat.
|
|
@returns True if the heartbeat notification is successfully sent.
|
Otherwise, False.
|
"""
|
logging.info('Sending a heartbeat')
|
|
event = cpcon.Heartbeat()
|
# Don't sent local timestamp for now.
|
data = event.SerializeToString()
|
try:
|
attributes = self._create_message_attributes(
|
cpcon.MSG_MOBLAB_HEARTBEAT)
|
message = self._create_message(data, attributes)
|
except ValueError:
|
logging.exception('Failed to create message.')
|
return False
|
return self._publish_notification(message)
|
|
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.
|
"""
|
logging.info('Send an event.')
|
if not event_type:
|
logging.info('Failed to send event without a type.')
|
return False
|
|
event = cpcon.RemoteEventMessage()
|
if event_data:
|
event.data = event_data
|
else:
|
event.data = ''
|
event.type = event_type
|
data = event.SerializeToString()
|
try:
|
attributes = self._create_message_attributes(
|
cpcon.MSG_MOBLAB_REMOTE_EVENT)
|
message = self._create_message(data, attributes)
|
except ValueError:
|
logging.exception('Failed to create message.')
|
return False
|
return self._publish_notification(message)
|
