def _BuildClusterInfo(self, cluster, host_infos):
"""Helper to build a ClusterInfo object from host messages.
Args:
cluster: a cluster entity
host_infos: a list of HostInfo messages.
Returns:
a ClusterInfo object.
"""
run_targets = set()
for host in host_infos:
run_targets.update([d.run_target
for d in host.device_infos if d.run_target])
run_target_messages = [api_messages.RunTarget(name=r) for r in run_targets]
return api_messages.ClusterInfo(
cluster_id=cluster.cluster,
total_devices=cluster.total_devices,
offline_devices=cluster.offline_devices,
available_devices=cluster.available_devices,
allocated_devices=cluster.allocated_devices,
device_count_timestamp=cluster.device_count_timestamp,
host_infos=host_infos,
run_targets=run_target_messages,
host_update_state_summary=datastore_entities.ToMessage(
cluster.host_update_state_summary),
host_count_by_harness_version=api_messages.MapToKeyValuePairMessages(
cluster.host_count_by_harness_version),
host_update_state_summaries_by_version=[
datastore_entities.ToMessage(summary) for summary
in cluster.host_update_state_summaries_by_version
])
def GetHost(self, request):
"""Fetches the information and notes of a given hostname.
Args:
request: an API request.
Returns:
a HostInfo object.
Raises:
endpoints.NotFoundException: If the given host does not exist.
endpoint.BadRequestException: If request includes history info with
negative limit.
"""
hostname = request.hostname
host = device_manager.GetHost(hostname)
if not host:
raise endpoints.NotFoundException("Host %s does not exist." % hostname)
device_query = datastore_entities.DeviceInfo.query(ancestor=host.key)
if not request.include_hidden:
device_query = device_query.filter(
datastore_entities.DeviceInfo.hidden == False) devices = device_query.fetch()
host_update_state = ndb.Key(
datastore_entities.HostUpdateState, hostname).get()
host_info = datastore_entities.ToMessage(
host, devices=devices, host_update_state_entity=host_update_state)
# TODO: deprecate "include_notes".
if request.include_notes:
host_notes = (
datastore_entities.Note.query().filter(
datastore_entities.Note.type == common.NoteType.HOST_NOTE).filter(
datastore_entities.Note.hostname == hostname).order(
-datastore_entities.Note.timestamp))
host_info.notes = [
datastore_entities.ToMessage(note) for note in host_notes
]
if request.include_host_state_history:
history_states = None
limit = request.host_state_history_limit
try:
history_states = device_manager.GetHostStateHistory(
hostname, limit=limit)
except ValueError as err:
raise endpoints.BadRequestException(err)
host_state_history = [
datastore_entities.ToMessage(state) for state in history_states
]
host_info.state_history = host_state_history
return host_info
def BatchGetLatestNotesByDevice(self, request):
"""Batch get notes of a device.
Args:
request: an API request.
Request Params:
device_serial: string, the serial of a lab device.
ids: a list of strings, the ids of notes to batch get.
Returns:
an api_messages.NoteCollection object.
"""
note_entities = []
for device_serial in request.device_serials:
query = (
datastore_entities.Note.query()
.filter(datastore_entities.Note.type == common.NoteType.DEVICE_NOTE)
.filter(datastore_entities.Note.device_serial == device_serial)
.order(-datastore_entities.Note.timestamp))
note_entities += list(query.fetch(1))
note_msgs = [
datastore_entities.ToMessage(entity) for entity in note_entities
]
return api_messages.NoteCollection(
notes=note_msgs, more=False, next_cursor=None, prev_cursor=None)
def NewNote(self, request):
"""Submits a note for this device.
Args:
request: an API request.
Returns:
an api_messages.Note object.
"""
timestamp = request.timestamp
# Datastore only accepts UTC times. Doing a conversion if necessary.
if timestamp.utcoffset() is not None:
timestamp = timestamp.replace(tzinfo=None) - timestamp.utcoffset()
note = datastore_entities.Note(
type=common.NoteType.DEVICE_NOTE,
hostname=request.hostname,
device_serial=request.device_serial,
user=request.user,
timestamp=timestamp,
message=request.message,
offline_reason=request.offline_reason,
recovery_action=request.recovery_action)
note.put()
return datastore_entities.ToMessage(note)
def testNoteFromEntity(self):
"""Tests converting a Note datastore entity a corresponding message."""
note_entity = datastore_entities.Note(
user='******',
timestamp=TIMESTAMP,
message='Hello, World',
offline_reason='something reasonable',
recovery_action='press the button',
type=common.NoteType.UNKNOWN,
cluster_id='acluser',
hostname='ahost',
device_serial='adevice',
event_time=TIMESTAMP)
note_message = datastore_entities.ToMessage(note_entity)
self.assertIsNone(note_message.id)
self.assertEqual(note_entity.user, note_message.user)
self.assertEqual(note_entity.timestamp, note_message.timestamp)
self.assertEqual(note_entity.message, note_message.message)
self.assertEqual(note_entity.offline_reason,
note_message.offline_reason)
self.assertEqual(note_entity.recovery_action,
note_message.recovery_action)
self.assertEqual(note_entity.type, note_message.type)
self.assertEqual(note_entity.cluster_id, note_message.cluster_id)
self.assertEqual(note_entity.hostname, note_message.hostname)
self.assertEqual(note_entity.device_serial, note_message.device_serial)
self.assertEqual(note_entity.event_time, note_message.event_time)
def GetDeviceSnapshot(self, request):
"""Gets a snapshot of all the devices and their properties for a given date.
If no date is provided, it will return a snapshot of the current state of
devices that would be reported.
Args:
request: an API request.
Returns:
a DeviceSnapshot object.
"""
logging.info('GetDevicesSnapshot request: %s', request)
date = None
if request.date:
date = dateutil.parser.parse(request.date).date()
device_snapshot = device_info_reporter.GetDeviceSnapshotForDate(date)
snapshot = DeviceSnapshot()
snapshot.updateTime = device_snapshot.update_time
# Message class does not support datetime.date so it has to be converted
# to a datetime.datetime object
report_date = datetime.datetime.combine(device_snapshot.date,
ZERO_TIME)
snapshot.date = report_date
snapshot.devices = [
datastore_entities.ToMessage(device_snapshot)
for device_snapshot in device_snapshot.devices
]
return snapshot
def testLabInfoFromEntity(self):
"""Test converting from lab_info to lab_info message."""
key = ndb.Key(datastore_entities.LabInfo, 'alab')
lab_info = datastore_entities.LabInfo(
key=key,
lab_name='alab',
update_timestamp=TIMESTAMP,
host_update_state_summary=datastore_entities.
HostUpdateStateSummary(total=2),
host_update_state_summaries_by_version=[
datastore_entities.HostUpdateStateSummary(total=2,
target_version='v1')
])
lab_config = datastore_entities.LabConfig(
id='alab', owners=['user1', 'user2', 'user3'])
msg = datastore_entities.ToMessage(lab_info, lab_config)
self.assertEqual(lab_info.lab_name, msg.lab_name)
self.assertEqual(lab_config.owners, msg.owners)
self.assertEqual(lab_info.update_timestamp, msg.update_timestamp)
self.assertEqual(lab_info.host_update_state_summary.total,
msg.host_update_state_summary.total)
self.assertEqual(
lab_info.host_update_state_summaries_by_version[0].total,
msg.host_update_state_summaries_by_version[0].total)
self.assertEqual(
lab_info.host_update_state_summaries_by_version[0].target_version,
msg.host_update_state_summaries_by_version[0].target_version)
def ListNotes(self, request):
"""List notes of a host.
Args:
request: an API request.
Returns:
an api_messages.NoteCollection object.
"""
query = (
datastore_entities.Note.query()
.filter(datastore_entities.Note.hostname == request.hostname)
.order(-datastore_entities.Note.timestamp))
if not request.include_device_notes:
query = query.filter(
datastore_entities.Note.type == common.NoteType.HOST_NOTE)
note_entities, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, request.cursor, backwards=request.backwards)
note_msgs = [
datastore_entities.ToMessage(entity) for entity in note_entities
]
return api_messages.NoteCollection(
notes=note_msgs,
more=bool(next_cursor),
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def GetCluster(self, request):
"""Fetches the information/status for a given cluster id.
Args:
request: an API request.
Returns:
a ClusterInfo message.
Raises:
endpoints.BadRequestException: If the given cluster ID is invalid.
"""
cluster_id = request.cluster_id
cluster = device_manager.GetCluster(cluster_id)
if not cluster:
raise endpoints.NotFoundException(
"Cluster [%s] does not exist." % cluster_id)
host_msgs = []
if request.include_hosts:
host_msgs = self._GetHostsForCluster(cluster_id)
cluster_info = self._BuildClusterInfo(cluster, host_msgs)
if request.include_notes:
cluster_notes = datastore_entities.ClusterNote.query()
cluster_notes = cluster_notes.filter(
datastore_entities.ClusterNote.cluster == cluster_id)
notes = [datastore_entities.ToMessage(n.note)
for n in cluster_notes.iter()]
cluster_info.notes = sorted(
notes, key=lambda x: x.timestamp, reverse=True)
return cluster_info
def ListCommandAttempts(self, api_request):
"""Get command attempts satisfy the condition.
Args:
api_request: api request may contain a hostname or device serial
Returns:
collection of command attempts
"""
query = datastore_entities.CommandAttempt.query(
namespace=common.NAMESPACE)
if api_request.hostname is not None:
query = query.filter(datastore_entities.CommandAttempt.hostname ==
api_request.hostname)
if api_request.device_serial is not None:
query = query.filter(datastore_entities.CommandAttempt.
device_serial == api_request.device_serial)
query = query.order(-datastore_entities.CommandAttempt.create_time)
if api_request.offset is not None and api_request.count is not None:
offset = api_request.offset
count = api_request.count
else:
offset = 0
count = common.DEFAULT_ROW_COUNT
command_attempt_entities = query.fetch(count, offset=offset)
attempts = [
datastore_entities.ToMessage(attempt)
for attempt in command_attempt_entities
]
return api_messages.CommandAttemptMessageCollection(
command_attempts=attempts)
def ListTestHarnessImages(self, request):
"""List test harness image metadata.
Args:
request: an API request.
Returns:
An api_messages.TestHarnessImageMetadataCollection object.
"""
query = datastore_entities.TestHarnessImageMetadata.query()
if request.tag_prefix:
query = (
query.filter(datastore_entities.TestHarnessImageMetadata.current_tags
>= request.tag_prefix).filter(
datastore_entities.TestHarnessImageMetadata
.current_tags < request.tag_prefix + u"\ufffd")
.order(-datastore_entities.TestHarnessImageMetadata.current_tags))
query = query.order(
-datastore_entities.TestHarnessImageMetadata.create_time)
entities, _, next_cursor = datastore_util.FetchPage(query, request.count,
request.cursor)
image_msgs = [datastore_entities.ToMessage(entity) for entity in entities]
return api_messages.TestHarnessImageMetadataCollection(
images=image_msgs, next_cursor=next_cursor)
def testStateHistoryFromEntity(self):
"""Tests converting a DeviceStateHistory datastore entity a message."""
entity = datastore_entities.DeviceStateHistory(device_serial='a1',
timestamp=TIMESTAMP,
state='Gone')
message = datastore_entities.ToMessage(entity)
self.assertEqual(entity.timestamp, message.timestamp)
self.assertEqual(entity.state, message.state)
def testDeviceBlocklistToMessage(self):
blocklist = datastore_test_util.CreateDeviceBlocklist('alab', 'auser')
msg = datastore_entities.ToMessage(blocklist)
self.assertIsNotNone(msg.key_id)
self.assertIsNotNone(msg.create_timestamp)
self.assertEqual('alab', msg.lab_name)
self.assertEqual('lab outage', msg.note)
self.assertEqual('auser', msg.user)
def GetDeviceBlocklist(self, request):
"""Get a device blocklist.
Args:
request: API request with key id.
Returns:
a DeviceBlocklistMessage object.
"""
blocklist = datastore_entities.DeviceBlocklist.get_by_id(
request.key_id)
return datastore_entities.ToMessage(blocklist)
def NewDeviceBlocklist(self, request):
"""Create a new device blocklist.
Args:
request: API request that includes device blocklist information.
Returns:
a DeviceBlocklistMessage object.
"""
blocklist = datastore_entities.DeviceBlocklist(
lab_name=request.lab_name, note=request.note, user=request.user)
blocklist.put()
return datastore_entities.ToMessage(blocklist)
def GetDevice(self, request):
"""Fetches the information and notes of a given device.
Args:
request: an API request.
Returns:
a DeviceInfo object.
Raises:
endpoints.NotFoundException: If the given device does not exist.
"""
device_serial = request.device_serial
device = device_manager.GetDevice(
hostname=request.hostname, device_serial=device_serial)
if not device:
raise endpoints.NotFoundException(
"Device {0} does not exist.".format(device_serial))
device_info = datastore_entities.ToMessage(device)
# TODO: deprecate "include_notes".
if request.include_notes:
device_notes = (
datastore_entities.Note.query().filter(
datastore_entities.Note.type == common.NoteType.DEVICE_NOTE)
.filter(datastore_entities.Note.device_serial == device_serial).order(
-datastore_entities.Note.timestamp))
device_info.notes = [
datastore_entities.ToMessage(note) for note in device_notes
]
if request.include_history:
histories = device_manager.GetDeviceStateHistory(device.hostname,
device_serial)
device_info.history = [datastore_entities.ToMessage(h) for h in histories]
if request.include_utilization:
utilization = device_manager.CalculateDeviceUtilization(device_serial)
device_info.utilization = utilization
return device_info
def _PublishHostMessage(hostname):
"""Publish host message to pubsub."""
if not env_config.CONFIG.use_google_api:
logging.warning(
'Unabled to send host message to pubsub: use_google_api=False')
return
host = device_manager.GetHost(hostname)
devices = device_manager.GetDevicesOnHost(hostname)
host_message = datastore_entities.ToMessage(host)
host_message.device_infos = [datastore_entities.ToMessage(d) for d in devices]
encoded_message = protojson.encode_message(host_message) # pytype: disable=module-attr
# TODO: find a better way to add event publish timestamp.
msg_dict = json.loads(encoded_message)
msg_dict['publish_timestamp'] = _Now().isoformat()
data = common.UrlSafeB64Encode(json.dumps(msg_dict))
_PubsubClient.PublishMessages(
HOST_AND_DEVICE_PUBSUB_TOPIC,
[{
'data': data,
'attributes': {
'type': 'host',
}
}])
def testDeviceBlocklistArchiveToMessage(self):
blocklist = datastore_test_util.CreateDeviceBlocklist('alab', 'auser')
blocklist_archive = (
datastore_entities.DeviceBlocklistArchive.FromDeviceBlocklist(
blocklist, 'another_user'))
blocklist_archive.put()
msg = datastore_entities.ToMessage(blocklist_archive)
self.assertEqual(msg.device_blocklist.create_timestamp,
msg.start_timestamp)
self.assertIsNotNone(msg.end_timestamp)
self.assertEqual('another_user', msg.archived_by)
self.assertEqual('alab', msg.device_blocklist.lab_name)
self.assertEqual('lab outage', msg.device_blocklist.note)
self.assertEqual('auser', msg.device_blocklist.user)
def testNoteFromEntity_withId(self):
note_entity = datastore_entities.Note(
key=ndb.Key(datastore_entities.Note, 123456789),
message='Hello, World',
offline_reason='something reasonable',
recovery_action='press the button',
type=common.NoteType.UNKNOWN)
note_message = datastore_entities.ToMessage(note_entity)
self.assertEqual('123456789', note_message.id)
self.assertEqual(note_entity.message, note_message.message)
self.assertEqual(note_entity.offline_reason,
note_message.offline_reason)
self.assertEqual(note_entity.recovery_action,
note_message.recovery_action)
self.assertEqual(note_entity.type, note_message.type)
def testPredefinedMessageFromEntity(self):
entity = datastore_entities.PredefinedMessage(
key=ndb.Key(datastore_entities.PredefinedMessage, 123456789),
lab_name='lab-name-01',
type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
content='device offline reason 1',
create_timestamp=TIMESTAMP,
used_count=4)
msg = datastore_entities.ToMessage(entity)
self.assertEqual(123456789, msg.id)
self.assertEqual(entity.lab_name, msg.lab_name)
self.assertEqual(entity.type, msg.type)
self.assertEqual(entity.content, msg.content)
self.assertEqual(entity.create_timestamp, msg.create_timestamp)
self.assertEqual(entity.used_count, msg.used_count)
def GetMetadata(self, request):
"""Get a host metadata.
Args:
request: an API request.
Returns:
an api_messages.HostMetadata object.
"""
metadata = datastore_entities.HostMetadata.get_by_id(request.hostname)
if not metadata:
metadata = datastore_entities.HostMetadata(hostname=request.hostname)
metadata_msg = datastore_entities.ToMessage(metadata)
return metadata_msg
def testRequestFromEntity(self):
"""Tests converting a Request entity to a message."""
request = datastore_entities.Request(
id='id',
user='******',
command_infos=[
datastore_entities.CommandInfo(command_line='command_line',
shard_count=2,
run_count=3,
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.UNKNOWN)
message = datastore_entities.ToMessage(request)
self.AssertEqualRequest(request, message)
def Restore(self, request):
"""Restore this host.
Args:
request: an API request.
Returns:
an updated HostInfo
Raises:
endpoints.NotFoundException: If the given device does not exist.
"""
host = device_manager.RestoreHost(request.hostname)
if not host:
raise endpoints.NotFoundException("Host %s does not exist." %
request.hostname)
return datastore_entities.ToMessage(host)
def testRequestFromEntity_canceledRequest(self):
"""Tests converting a cancelled Request entity to a message."""
request = datastore_entities.Request(
id='id',
user='******',
command_infos=[
datastore_entities.CommandInfo(command_line='command_line',
shard_count=2,
run_count=3,
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.CANCELED,
cancel_reason=common.CancelReason.QUEUE_TIMEOUT)
message = datastore_entities.ToMessage(request)
self.AssertEqualRequest(request, message)
def testHarnessImageMetadataToMessage(self):
entity = datastore_entities.TestHarnessImageMetadata(
repo_name='gcr.io/dockerized-tradefed/tradefed',
digest='sha1',
test_harness='tradefed',
test_harness_version='111111',
current_tags=['111111'],
create_time=TIMESTAMP_OLD,
sync_time=TIMESTAMP_NEW)
msg = datastore_entities.ToMessage(entity)
self.assertEqual('gcr.io/dockerized-tradefed/tradefed', msg.repo_name)
self.assertEqual('sha1', msg.digest)
self.assertEqual('tradefed', msg.test_harness)
self.assertEqual('111111', msg.test_harness_version)
self.assertEqual(['111111'], msg.tags)
self.assertEqual(TIMESTAMP_OLD, msg.create_time)
def _ListLabs(self, request):
"""ListLabs without owner filter. Some labs don't have config."""
query = datastore_entities.LabInfo.query()
query = query.order(datastore_entities.LabInfo.key)
labs, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, page_cursor=request.cursor)
lab_config_keys = [
ndb.Key(datastore_entities.LabConfig, lab.lab_name) for lab in labs]
lab_configs = ndb.get_multi(lab_config_keys)
lab_infos = [datastore_entities.ToMessage(lab, lab_config)
for lab, lab_config in zip(labs, lab_configs)]
return api_messages.LabInfoCollection(
lab_infos=lab_infos,
more=bool(next_cursor),
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def testMessageFromEntity_cancelRequest_invalidRequest(self):
"""Tests converting a Request object to a message."""
request_key = ndb.Key(datastore_entities.Request, '123')
request = datastore_entities.Request(
key=request_key,
user='******',
command_infos=[
datastore_entities.CommandInfo(command_line='command_line',
shard_count=2,
run_count=3,
cluster='cluster',
run_target='run_target')
],
state=common.RequestState.UNKNOWN,
cancel_reason=common.CancelReason.INVALID_REQUEST)
message = datastore_entities.ToMessage(request)
self.AssertEqualRequest(request, message)
def _GetHostsForCluster(self, cluster_id):
"""Get hosts and their devices for a cluster.
Args:
cluster_id: cluster id
Returns:
a list of HostInfoMessages include devices.
"""
hosts = (datastore_entities.HostInfo.query()
.filter(datastore_entities.HostInfo.clusters == cluster_id)
.filter(datastore_entities.HostInfo.hidden == False) .fetch())
host_msgs = []
for host in hosts:
devices = (datastore_entities.DeviceInfo.query(ancestor=host.key)
.filter(datastore_entities.DeviceInfo.hidden == False) .fetch())
host_msgs.append(datastore_entities.ToMessage(host, devices))
return host_msgs
def testHostNoteFromEntity(self):
entity = datastore_entities.HostNote(
key=ndb.Key(datastore_entities.HostNote, 123456789),
hostname='hostname_1',
note=datastore_entities.Note(user='******',
timestamp=TIMESTAMP,
offline_reason='offline_reason_1',
recovery_action='recovery_action_1',
message='message_1'))
msg = datastore_entities.ToMessage(entity)
self.assertEqual('123456789', msg.id)
self.assertEqual(entity.hostname, msg.hostname)
self.assertEqual(entity.note.user, msg.user)
self.assertEqual(entity.note.timestamp, msg.update_timestamp)
self.assertEqual(entity.note.offline_reason, msg.offline_reason)
self.assertEqual(entity.note.recovery_action, msg.recovery_action)
self.assertEqual(entity.note.message, msg.message)
def testDeviceInfoFromEntity(self):
"""Test converting from device_info to device_info message."""
entity = self._CreateMockDeviceInfoEntity()
msg = datastore_entities.ToMessage(entity)
self.assertEqual(entity.device_serial, msg.device_serial)
self.assertEqual(entity.hostname, msg.hostname)
self.assertEqual('Available', msg.state)
self.assertEqual('alab', msg.lab_name)
self.assertEqual('atestharness', msg.test_harness)
self.assertEqual('acluster', msg.cluster)
self.assertEqual('atp-us-mtv-43', msg.host_group)
self.assertEqual(['apct', 'asit'], msg.pools)
self.assertEqual(
entity.extra_info,
api_messages.KeyValuePairMessagesToMap(msg.extra_info))
self.assertEqual(common.RecoveryState.FIXED, entity.recovery_state)
self.assertEqual(TIMESTAMP, msg.timestamp)
self.assertEqual(TIMESTAMP, msg.last_recovery_time)