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 CreateDeviceCountSummary(run_target, offline=0, available=0, allocated=0):
"""Create a device count summary."""
return datastore_entities.DeviceCountSummary(run_target=run_target,
total=offline + available +
allocated,
offline=offline,
available=available,
allocated=allocated)
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 _DoCountDeviceForHost(host, devices):
"""Actually count devices for a host."""
if not host:
return
if not devices:
logging.info("No devices reported for host [%s]", host.hostname)
if not host.total_devices and not host.device_count_summaries:
return
# If there is no devices but the total_devices is not 0, we need to clear
# the count.
now = common.Now()
host.total_devices = 0
host.offline_devices = 0
host.available_devices = 0
host.allocated_devices = 0
device_counts = {}
for device in devices or []:
device_count = device_counts.get(device.run_target)
if not device_count:
device_count = datastore_entities.DeviceCountSummary(
run_target=device.run_target,
timestamp=now)
device_counts[device.run_target] = device_count
device_count.total += 1
host.total_devices += 1
if device.state in common.DEVICE_AVAILABLE_STATES:
host.available_devices += 1
device_count.available += 1
elif device.state in common.DEVICE_ALLOCATED_STATES:
host.allocated_devices += 1
device_count.allocated += 1
else:
host.offline_devices += 1
device_count.offline += 1
host.device_count_timestamp = now
host.device_count_summaries = list(device_counts.values())