def testLabInfo(self): key = ndb.Key(datastore_entities.LabInfo, 'alab') lab_info = datastore_entities.LabInfo( key=key, lab_name='alab', update_timestamp=TIMESTAMP_NEW, host_update_state_summary=datastore_entities. HostUpdateStateSummary(total=2)) lab_info.put() lab_info_res = key.get() self.assertEqual('alab', lab_info_res.lab_name) self.assertIsNotNone(lab_info_res.update_timestamp) self.assertEqual(2, lab_info_res.host_update_state_summary.total)
def SyncCommandAttempt(request_id, command_id, attempt_id): """Sync the command attempt. Reset (error) the attempt if running but inactive, otherwise re-add to the sync queue to check back again later. Args: request_id: Request ID for the command attempt to sync command_id: Command ID for the command attempt to sync attempt_id: Attempt ID for the command attempt to sync """ now = Now() attempt_key = ndb.Key(datastore_entities.Request, request_id, datastore_entities.Command, command_id, datastore_entities.CommandAttempt, attempt_id, namespace=common.NAMESPACE) attempt = attempt_key.get() if not attempt: logging.warning( 'No attempt found to sync. Request %s Command %s Attempt %s', request_id, command_id, attempt_id) return if attempt.state in common.FINAL_COMMAND_STATES: # No need to sync attempt in final states logging.info( 'Attempt reached final state %s. Request %s Command %s Attempt %s', attempt.state, request_id, command_id, attempt_id) return inactive_time = now - attempt.update_time if inactive_time > datetime.timedelta( minutes=command_manager.MAX_COMMAND_EVENT_DELAY_MIN): logging.info( 'Resetting command task %s which has been inactive for %s.', attempt.task_id, inactive_time) event = command_event.CommandEvent( task_id=attempt.task_id, attempt_id=attempt_id, type=common.InvocationEventType.INVOCATION_COMPLETED, data={'error': common.TASK_RESET_ERROR_MESSAGE}, time=now) commander.ProcessCommandEvent(event) return # Add to the sync queue to check again later. command_manager.AddToSyncCommandAttemptQueue(attempt)
def ConvertToKey(model_cls, id_): """Convert to a ndb.Key object. Args: model_cls: a model class. id_: an object ID. Returns: a ndb.Key object. """ try: id_ = int(id_) except ValueError: pass return ndb.Key(model_cls, id_)
def _CreateMockTestRunSequence( self, sequence_id='sequence_id', configs=None, state=ndb_models.TestRunSequenceState.RUNNING): """Create a mock ndb_models.TestRunSequence object.""" sequence = ndb_models.TestRunSequence( state=state, test_run_configs=configs or [], finished_test_run_ids=[], ) sequence.key = ndb.Key(ndb_models.TestRunSequence, sequence_id) sequence.put() return sequence
def GetHostUpdateStateHistories(hostname, limit=DEFAULT_HOST_UPDATE_STATE_HISTORY_SIZE): """Function to get host update state history from NDB. Args: hostname: host name. limit: an integer about the max number of state history returned. Returns: A datastore entity list of HostUpdateStateHistory. """ return (datastore_entities.HostUpdateStateHistory .query(ancestor=ndb.Key(datastore_entities.HostInfo, hostname)) .order(-datastore_entities.HostUpdateStateHistory.update_timestamp) .fetch(limit=limit))
def GetCommandAttempt(request_id, command_id, attempt_id): """Returns the command attempt that matches the attempt id. Args: request_id: id for the parent request command_id: id for the parent command attempt_id: id for the attempt Returns: a CommandAttempt entity """ return ndb.Key( datastore_entities.Request, request_id, datastore_entities.Command, command_id, datastore_entities.CommandAttempt, attempt_id, namespace=common.NAMESPACE).get()
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 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 testGetOrCreatePredefinedMessage_ExistingMessage(self): lab_name = "alab" content = "content1" message_id = 111 datastore_entities.PredefinedMessage( key=ndb.Key(datastore_entities.PredefinedMessage, message_id), lab_name=lab_name, type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, content=content, used_count=2).put() message = note_manager.GetOrCreatePredefinedMessage( api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, lab_name, content) self.assertEqual(message_id, message.key.id())
def testFetchPage_backwardsWithAncestorQuery(self): self.ndb_host_0 = datastore_test_util.CreateHost( cluster='free', hostname='host_0', timestamp=self.TIMESTAMP, host_state=api_messages.HostState.RUNNING, ) self.ndb_host_1 = datastore_test_util.CreateHost( cluster='paid', hostname='host_1', timestamp=self.TIMESTAMP, device_count_timestamp=self.TIMESTAMP, ) self._CreateHostInfoHistory(self.ndb_host_1).put() self.ndb_host_1.host_state = api_messages.HostState.UNKNOWN self.ndb_host_1.timestamp += datetime.timedelta(hours=1) self._CreateHostInfoHistory(self.ndb_host_1).put() self.ndb_host_1.host_state = api_messages.HostState.GONE self.ndb_host_1.timestamp += datetime.timedelta(hours=1) self._CreateHostInfoHistory(self.ndb_host_1).put() self._CreateHostInfoHistory(self.ndb_host_0).put() self.ndb_host_0.host_state = api_messages.HostState.KILLING self.ndb_host_0.timestamp += datetime.timedelta(hours=1) self._CreateHostInfoHistory(self.ndb_host_0).put() self.ndb_host_0.host_state = api_messages.HostState.GONE self.ndb_host_0.timestamp += datetime.timedelta(hours=1) self._CreateHostInfoHistory(self.ndb_host_0).put() # First page query = (datastore_entities.HostInfoHistory.query(ancestor=ndb.Key( datastore_entities.HostInfo, self.ndb_host_0.hostname)).order( -datastore_entities.HostInfoHistory.timestamp)) histories, prev_cursor, next_cursor = datastore_util.FetchPage( query, 2) self.assertEqual(2, len(histories)) self.assertIsNone(prev_cursor) self.assertIsNotNone(next_cursor) # Back to first page (ancestor query with backwards) histories, prev_cursor, next_cursor = datastore_util.FetchPage( query, 2, backwards=True) self.assertEqual(2, len(histories)) self.assertEqual(self.ndb_host_0.hostname, histories[0].hostname) self.assertEqual(api_messages.HostState.GONE, histories[0].host_state) self.assertEqual(self.ndb_host_0.hostname, histories[1].hostname) self.assertEqual(api_messages.HostState.KILLING, histories[1].host_state)
def testClusterNote(self): cluster_note = datastore_entities.ClusterNote(cluster='free') key = ndb.Key(datastore_entities.Note, datastore_entities.Note.allocate_ids(1)[0].id()) note = datastore_entities.Note(key=key, user='******', timestamp=TIMESTAMP_OLD, message='Hello, World') cluster_note.note = note cluster_note_key = cluster_note.put() queried_cluster_note = cluster_note_key.get() self.assertEqual('free', queried_cluster_note.cluster) self.assertEqual(cluster_note_key, queried_cluster_note.key) self.assertEqual(note.message, queried_cluster_note.note.message) self.assertEqual(note.timestamp, queried_cluster_note.note.timestamp) self.assertEqual(note.user, queried_cluster_note.note.user)
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 testCreateTestRun_withRerunConfigs(self): test = self._CreateMockTest() config1 = self._CreateMockTestRunConfig(test) config2 = self._CreateMockTestRunConfig(test, 'rt1;rt2') config3 = self._CreateMockTestRunConfig(test, 't1;t2;t3') test_run = test_kicker.CreateTestRun( ['label'], config1, rerun_configs=[config2, config3]) sequence_id = test_run.sequence_id sequence = ndb.Key(ndb_models.TestRunSequence, sequence_id).get() self.assertEqual(len(sequence.test_run_configs), 3) self.assertModelEqual(sequence.test_run_configs[0], config1) self.assertModelEqual(sequence.test_run_configs[1], config2) self.assertModelEqual(sequence.test_run_configs[2], config3)
def SetTestEnvironment(request_id, test_env): """Sets a test environment for a request. Args: request_id: a request ID. test_env: a datastore_entities.TestEnvironment object. """ request_key = ndb.Key( datastore_entities.Request, request_id, namespace=common.NAMESPACE) test_env_to_put = datastore_entities.TestEnvironment.query( ancestor=request_key).get() if not test_env_to_put: test_env_to_put = datastore_entities.TestEnvironment(parent=request_key) test_env_to_put.populate(**test_env.to_dict()) test_env_to_put.put()
def GetHostStateHistory(hostname, limit=DEFAULT_HOST_HISTORY_SIZE): """Function to get host state history from NDB. Args: hostname: host name. limit: an integer about the max number of state history returned. Returns: a list of host state history. """ host_key = ndb.Key(datastore_entities.HostInfo, hostname) if limit < 0 or limit > MAX_HOST_HISTORY_SIZE: raise ValueError("size of host state history should be in range 0 to %d," "but got %d" % MAX_HOST_HISTORY_SIZE % limit) return (datastore_entities.HostStateHistory.query(ancestor=host_key) .order(-datastore_entities.HostStateHistory.timestamp) .fetch(limit=limit))
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 CreateRequest(user, command_infos, priority=None, queue_timeout_seconds=None, type_=None, request_id=None, plugin_data=None, max_retry_on_test_failures=None, prev_test_context=None, max_concurrent_tasks=None): """Create a new request and add it to the request_queue. Args: user: a requesting user. command_infos: a list of CommandInfo entities. priority: a request priority. queue_timeout_seconds: a request timeout in seconds. type_: a request type. request_id: request_id, used for easy testing. plugin_data: a map that contains the plugin data. max_retry_on_test_failures: the max number of completed but failed attempts for each command. prev_test_context: a previous test context. max_concurrent_tasks: the max number of concurrent tasks at any given time. Returns: a Request entity, read only. """ if not request_id: request_id = _CreateRequestId() key = ndb.Key( datastore_entities.Request, request_id, namespace=common.NAMESPACE) request = datastore_entities.Request( key=key, user=user, command_infos=command_infos, priority=priority, queue_timeout_seconds=queue_timeout_seconds, type=type_, plugin_data=plugin_data, max_retry_on_test_failures=max_retry_on_test_failures, prev_test_context=prev_test_context, max_concurrent_tasks=max_concurrent_tasks, state=common.RequestState.UNKNOWN) request.put() return request
def testProcessCommandEvent_withRequestSync(self, mock_store_event, mock_event_legacy_count): _, request_id, _, command_id = self.command.key.flat() sync_key = ndb.Key( datastore_entities.RequestSyncStatus, request_id, namespace=common.NAMESPACE) sync_status = datastore_entities.RequestSyncStatus( key=sync_key, request_id=request_id) sync_status.put() event = command_event_test_util.CreateTestCommandEvent( request_id, command_id, "0000000", "InvocationCompleted") command_event_handler.ProcessCommandEvent(event) mock_store_event.assert_called_once_with(event) mock_event_legacy_count.assert_not_called()
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 _CountDeviceForHost(hostname): """Count devices for a host. Args: hostname: the host name to count. """ host = GetHost(hostname) if not host: return devices = ( datastore_entities.DeviceInfo .query(ancestor=ndb.Key(datastore_entities.HostInfo, hostname)) .filter(datastore_entities.DeviceInfo.hidden == False) .fetch(projection=[ datastore_entities.DeviceInfo.run_target, datastore_entities.DeviceInfo.state])) _DoCountDeviceForHost(host, devices) host.put()
def GetDeviceStateHistory(hostname, device_serial): """Retrieve a device's state history. Limit to MAX_HISTORY_SIZE Args: hostname: hostname device_serial: a device serial. Returns: a list of DeviceStateHistory entities. """ device_key = ndb.Key( datastore_entities.HostInfo, hostname, datastore_entities.DeviceInfo, device_serial) return (datastore_entities.DeviceStateHistory.query(ancestor=device_key) .order(-datastore_entities.DeviceStateHistory.timestamp) .fetch(limit=MAX_DEVICE_HISTORY_SIZE))
def GetDevice(hostname=None, device_serial=None): """Retrieve a device given a device serial and its hostname. Args: hostname: hostname device_serial: a device serial. Returns: The device information corresponding to the given device serial. """ if hostname: return ndb.Key( datastore_entities.HostInfo, hostname, datastore_entities.DeviceInfo, device_serial).get() return (datastore_entities.DeviceInfo.query() .filter( datastore_entities.DeviceInfo.device_serial == device_serial) .order(-datastore_entities.DeviceInfo.timestamp).get())
def testHostInfo(self): key = ndb.Key(datastore_entities.HostInfo, 'ahost') host_info = datastore_entities.HostInfo( key=key, hostname='ahost', host_state=api_messages.HostState.RUNNING, device_count_summaries=[ datastore_entities.DeviceCountSummary(run_target='r1', total=1, available=1) ]) host_info.put() host_info_res = key.get() self.assertEqual('ahost', host_info_res.hostname) self.assertEqual(api_messages.HostState.RUNNING, host_info_res.host_state) self.assertFalse(host_info_res.is_bad) self.assertIsNotNone(host_info.timestamp)
def testLog_entity(self): """Tests that log entries can be written and retrieved per entity.""" test_run = ndb_models.TestRun(id='test_run_id') # Actual entity system_key = ndb.Key('System', 1) # Fake 'system' entity event_log.Info(test_run, 'entity message') event_log.Info(system_key, 'system message') # Can retrieve entity log entries entity_entries = event_log.GetEntries(test_run) self.assertLen(entity_entries, 1) self._VerifyEntry(entity_entries[0], ndb_models.EventLogLevel.INFO, 'entity message') # Can retrieve system log entries system_entries = event_log.GetEntries(system_key) self.assertLen(system_entries, 1) self._VerifyEntry(system_entries[0], ndb_models.EventLogLevel.INFO, 'system message')
def testHostInfo_gone(self): key = ndb.Key(datastore_entities.HostInfo, 'ahost') host_info = datastore_entities.HostInfo( key=key, hostname='ahost', host_state=api_messages.HostState.GONE, timestamp=TIMESTAMP_NEW, device_count_summaries=[ datastore_entities.DeviceCountSummary(run_target='r1', total=1, available=1) ]) host_info.put() host_info_res = key.get() self.assertEqual('ahost', host_info_res.hostname) self.assertEqual(TIMESTAMP_NEW, host_info_res.timestamp.replace(tzinfo=None)) self.assertEqual(api_messages.HostState.GONE, host_info_res.host_state) self.assertTrue(host_info_res.is_bad)
def testPreparePredefinedMessageForNote_withValidId(self): message_id = 111 lab_name = "alab" content = "content1" datastore_entities.PredefinedMessage( key=ndb.Key(datastore_entities.PredefinedMessage, message_id), lab_name=lab_name, type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, content=content, used_count=2).put() message = note_manager.PreparePredefinedMessageForNote( api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, message_id=message_id) self.assertEqual(message_id, message.key.id()) self.assertEqual(lab_name, message.lab_name) self.assertEqual(content, message.content) self.assertEqual(3, message.used_count) # the used_count increases
def _UpdateHostUpdateStateWithEvent( event, target_version=common.UNKNOWN_TEST_HARNESS_VERSION): """Update the host with a host update state change event. Args: event: HostEvent object. target_version: The test harness version which the host updates to. """ entities_to_update = [] host_update_state_enum = api_messages.HostUpdateState(event.host_update_state) host_update_state = datastore_entities.HostUpdateState.get_by_id( event.hostname) if not host_update_state: host_update_state = datastore_entities.HostUpdateState( id=event.hostname, hostname=event.hostname) if (host_update_state.update_timestamp and event.timestamp and host_update_state.update_timestamp > event.timestamp): logging.info("Ignore outdated event.") else: host_update_state.populate( state=host_update_state_enum, update_timestamp=event.timestamp, update_task_id=event.host_update_task_id, display_message=event.host_update_state_display_message, target_version=target_version) entities_to_update.append(host_update_state) host_update_state_history = datastore_entities.HostUpdateStateHistory( parent=ndb.Key(datastore_entities.HostInfo, event.hostname), hostname=event.hostname, state=host_update_state_enum, update_timestamp=event.timestamp, update_task_id=event.host_update_task_id, display_message=event.host_update_state_display_message, target_version=target_version) entities_to_update.append(host_update_state_history) ndb.put_multi(entities_to_update)
def PreparePredefinedMessageForNote(message_type, message_id=None, lab_name=None, content=None, delta_count=1): """Prepare a PredefinedMessage to attach to a Note. This method prepares a PredefinedMessage in following ways: - if message_id is provided, find the message with id, or - if content is provide, get existing message matching the content, or create new message with the content - if neither is provided, return None Args: message_type: enum, common.PredefinedMessageType, type of PredefinedMessage. message_id: int, the ID of PredefinedMessage. lab_name: str, the lab where the message is created. content: str, content of the message. delta_count: the delta used_count to be added. Returns: An instance of datastore_entities.PredefinedMessage. Raises: InvalidParameterError: when the message_id is not valid or it leads to an wrong PredefineMessage type. """ predefined_message_entity = None if message_id: predefined_message_entity = ndb.Key( datastore_entities.PredefinedMessage, message_id).get() if (not predefined_message_entity or predefined_message_entity.type != message_type): raise InvalidParameterError("Invalid predefined_message_id: %s" % message_id) elif content: predefined_message_entity = GetOrCreatePredefinedMessage( message_type, lab_name, content) if predefined_message_entity: predefined_message_entity.used_count += delta_count return predefined_message_entity
def testCommandFromEntity(self): """Tests converting a Command entity to a message.""" command_key = ndb.Key(datastore_entities.Request, '1', datastore_entities.Command, '2') cmd = datastore_entities.Command( key=command_key, command_line='command_line2', cluster='cluster', run_target='run_target', run_count=10, state=common.CommandState.QUEUED, start_time=TIMESTAMP, end_time=None, create_time=TIMESTAMP, update_time=TIMESTAMP, cancel_reason=common.CancelReason.QUEUE_TIMEOUT, error_reason=common.ErrorReason.TOO_MANY_LOST_DEVICES, shard_count=2, shard_index=1) message = datastore_entities.ToMessage(cmd) self.AssertEqualCommand(cmd, message)
def CreateDevice(cluster, hostname, device_serial, lab_name=None, battery_level='100', hidden=False, device_type=api_messages.DeviceTypeMessage.PHYSICAL, timestamp=None, state='Available', product='product', run_target='run_target', next_cluster_ids=None, test_harness='tradefed', pools='pool_01', host_group='host_group_01', extra_info=None, last_recovery_time=None): """Create a device.""" ndb_device = datastore_entities.DeviceInfo( id=device_serial, parent=ndb.Key(datastore_entities.HostInfo, hostname), device_serial=device_serial, hostname=hostname, battery_level=battery_level, device_type=device_type, hidden=hidden, lab_name=lab_name, physical_cluster=cluster, clusters=[cluster] + (next_cluster_ids if next_cluster_ids else []), timestamp=timestamp, state=state, product=product, run_target=run_target, test_harness=test_harness, pools=[pools], host_group=host_group, extra_info=extra_info, last_recovery_time=last_recovery_time) ndb_device.put() return ndb_device