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
