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)