def _UpdateHostWithHostChangedEvent(event):
"""update the host with a host state changed event.
Args:
event: HostEvent object.
"""
host = GetHost(event.hostname)
if not host:
host = datastore_entities.HostInfo(id=event.hostname)
# For HostStateChangedEvent other than hostname, state and timestamp,
# everything else are optional.
host.hostname = event.hostname
host.lab_name = event.lab_name or host.lab_name or common.UNKNOWN_LAB_NAME
host.timestamp = event.timestamp
host.test_harness = event.test_harness or host.test_harness
host.test_harness_version = (
event.test_harness_version or host.test_harness_version)
host.extra_info = event.data or host.extra_info
host.hidden = False
host_state_history, host_history = _UpdateHostState(
host, event.host_state, event.timestamp)
entities_to_update = [host]
if host_state_history:
entities_to_update.append(host_state_history)
if host_history:
entities_to_update.append(host_history)
ndb.put_multi(entities_to_update)
return
def LeaseHostTasks(self, request):
"""Lease available command tasks for a given host.
Args:
request: a HostInfo
Returns:
a TaskList object.
"""
logging.debug("leasehosttasks: request=%s", request)
host = device_manager.GetHost(request.hostname)
if not host:
host = datastore_entities.HostInfo(
hostname=request.hostname,
host_group=request.cluster)
matcher = command_task_matcher.CommandTaskMatcher(request)
run_targets = matcher.GetRunTargets()
if not run_targets:
return CommandTaskList(tasks=[])
clusters = [request.cluster]
clusters.extend(request.next_cluster_ids)
leased_tasks = []
num_tasks = request.num_tasks
for cluster in clusters:
try:
cluster_leased_tasks = self._LeaseHostTasksForCluster(
matcher, cluster, host, num_tasks)
leased_tasks.extend(cluster_leased_tasks)
if num_tasks is not None:
num_tasks -= len(cluster_leased_tasks)
if num_tasks <= 0:
break
if matcher.IsEmpty():
break
except:
def _CreateMockHostInfoEntity(self):
"""Helper function to get mock host info entity."""
d1_count = datastore_entities.DeviceCountSummary(run_target='d1',
total=10,
offline=1,
available=5,
allocated=4,
timestamp=TIMESTAMP)
d2_count = datastore_entities.DeviceCountSummary(run_target='d2',
total=5,
offline=1,
available=3,
allocated=1,
timestamp=TIMESTAMP)
return datastore_entities.HostInfo(
hostname='hostname',
lab_name='alab',
host_group='atp-us-mtv-43',
physical_cluster='acluster',
pools=['apct', 'asit'],
host_state=api_messages.HostState.RUNNING,
assignee='auser',
extra_info={
'host_url': 'aurl',
},
device_count_summaries=[d1_count, d2_count],
last_recovery_time=TIMESTAMP,
recovery_state=common.RecoveryState.FIXED)
def CreateHost(cluster,
hostname,
lab_name=None,
hidden=False,
timestamp=None,
extra_info=None,
device_count_timestamp=None,
host_state=api_messages.HostState.UNKNOWN,
assignee=None,
device_count_summaries=None,
test_harness='TRADEFED',
test_harness_version='1234',
pools=None):
"""Create a host."""
total_devices = 0
offline_devices = 0
available_devices = 0
allocated_devices = 0
for c in device_count_summaries or []:
total_devices += c.total
offline_devices += c.offline
available_devices += c.available
allocated_devices += c.allocated
ndb_host = datastore_entities.HostInfo(
id=hostname,
hostname=hostname,
lab_name=lab_name,
physical_cluster=cluster,
host_group=cluster,
clusters=[cluster],
pools=pools or [],
timestamp=timestamp,
extra_info=extra_info,
hidden=hidden,
total_devices=total_devices,
offline_devices=offline_devices,
available_devices=available_devices,
allocated_devices=allocated_devices,
device_count_timestamp=device_count_timestamp,
host_state=host_state,
assignee=assignee,
device_count_summaries=device_count_summaries or [],
test_harness=test_harness,
test_harness_version=test_harness_version)
ndb_host.put()
return ndb_host
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 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 _UpdateHostWithDeviceSnapshotEvent(event):
"""update the host if the event is host info.
Update host state to RUNNING if the olds state is GONE.
Args:
event: HostEvent dictionary.
Returns:
a HostEntity.
"""
host = GetHost(event.hostname)
if not host:
host = datastore_entities.HostInfo(id=event.hostname)
host.hostname = event.hostname
host.lab_name = event.lab_name
# TODO: deprecate physical_cluster, use host_group.
host.physical_cluster = event.cluster_id
host.host_group = event.host_group
host.timestamp = event.timestamp
host.test_harness = event.test_harness
host.test_harness_version = event.test_harness_version
host.hidden = False
# TODO: deprecate clusters, use pools.
if event.cluster_id:
host.clusters = [event.cluster_id] + event.next_cluster_ids
else:
host.clusters = event.next_cluster_ids[:]
host.pools = event.pools
entities_to_update = [host]
if _IsNewTestHarnessInstance(host, event):
# If it's a new instance, we change the host state to RUNNING.
host_state_history, host_history = _UpdateHostState(
host, api_messages.HostState.RUNNING, event.timestamp)
if host_state_history:
entities_to_update.append(host_state_history)
if host_history:
entities_to_update.append(host_history)
# Extra info need to be update after checking _IsNewTestHarnessInstance,
# since we use insit_harness_start_time_ms in extra info.
host.extra_info = event.data
ndb.put_multi(entities_to_update)
return host