def _UpdateHostConfigs(host_config_pbs, cluster_config_pb, lab_config_pb):
"""Update host configs in HostInfo entities to ndb.
Args:
host_config_pbs: a list of host config protos.
cluster_config_pb: the cluster config proto.
lab_config_pb: the lab config proto.
"""
logging.debug('Updating host configs for <lab: %s, cluster: %s.>',
lab_config_pb.lab_name, cluster_config_pb.cluster_name)
host_config_keys = []
for host_config_pb in host_config_pbs:
host_config_keys.append(
ndb.Key(datastore_entities.HostConfig, host_config_pb.hostname))
entities = ndb.get_multi(host_config_keys)
entities_to_update = []
# Update the exist host config entity.
for entity, host_config_pb in zip(entities, host_config_pbs):
host_config_msg = lab_config_util.HostConfig(
host_config_pb, cluster_config_pb, lab_config_pb)
new_host_config_entity = datastore_entities.HostConfig.FromMessage(
host_config_msg)
if not _CheckConfigEntitiesEqual(entity, new_host_config_entity):
logging.debug('Updating host config entity: %s.', new_host_config_entity)
entities_to_update.append(new_host_config_entity)
ndb.put_multi(entities_to_update)
logging.debug('Host configs updated.')
def testGetPredefinedMessage_OK(self):
lab_name = "alab"
predefined_message_entities = [
datastore_entities.PredefinedMessage(
lab_name=lab_name,
type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
content="offline_reason1",
used_count=2),
datastore_entities.PredefinedMessage(
lab_name=lab_name,
type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
content="recovery_action1",
used_count=5),
]
ndb.put_multi(predefined_message_entities)
message_1 = note_manager.GetPredefinedMessage(
api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, lab_name,
"offline_reason1")
self.assertEqual("offline_reason1", message_1.content)
self.assertEqual(api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
message_1.type)
message_2 = note_manager.GetPredefinedMessage(
api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, lab_name,
"recovery_action1")
self.assertEqual("recovery_action1", message_2.content)
self.assertEqual(api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
message_2.type)
def _SetHostRecoveryState(hostname, recovery_state, assignee=None):
"""Set host's recovery state."""
host = GetHost(hostname)
if not host:
logging.error("Host %s doesn't exist.", hostname)
return
host.assignee = assignee
entities_to_update = []
if recovery_state == common.RecoveryState.VERIFIED:
host.recovery_state = common.RecoveryState.VERIFIED
host.assignee = host.assignee
host.last_recovery_time = common.Now()
host.timestamp = common.Now()
entities_to_update.append(_CreateHostInfoHistory(host))
devices = GetDevicesOnHost(hostname)
for device in devices:
if (device.recovery_state and
device.recovery_state != common.RecoveryState.UNKNOWN):
entities_to_update.extend(
_BuildDeviceRecoveryState(device, common.RecoveryState.VERIFIED))
# After it's verified, it goes into a state as no-one is recovering it.
host.recovery_state = common.RecoveryState.UNKNOWN
host.assignee = None
else:
host.recovery_state = recovery_state
host.last_recovery_time = common.Now()
host.timestamp = common.Now()
entities_to_update.append(_CreateHostInfoHistory(host))
entities_to_update.append(host)
ndb.put_multi(entities_to_update)
def UpdateGoneHost(hostname):
"""Set a host and its devices to GONE."""
logging.info("Set host %s and its devices to GONE.", hostname)
host = GetHost(hostname)
if host.host_state == api_messages.HostState.GONE:
logging.info("Host %s is already GONE.", hostname)
return
entities_to_update = []
now = common.Now()
host_state_history, host_history = _UpdateHostState(
host, api_messages.HostState.GONE, now)
entities_to_update.append(host)
if host_state_history:
entities_to_update.append(host_state_history)
if host_history:
entities_to_update.append(host_history)
devices = GetDevicesOnHost(hostname)
for device in devices or []:
if device.state == common.DeviceState.GONE:
continue
logging.debug("Set device %s to GONE.", device.device_serial)
device_state_history, device_history = _UpdateDeviceState(
device, common.DeviceState.GONE, now)
entities_to_update.append(device)
if device_state_history:
entities_to_update.append(device_state_history)
if device_history:
entities_to_update.append(device_history)
_DoCountDeviceForHost(host, devices)
ndb.put_multi(entities_to_update)
def testSyncInventoryGroupsToNdbHostGroupConfig(self):
ndb.put_multi([
datastore_entities.HostGroupConfig(
id='foo_jump',
lab_name='foo',
parent_groups=['foo_all', 'foo_bar']),
datastore_entities.HostGroupConfig(id='foo_bar',
lab_name='foo',
parent_groups=[
'foo_all',
])
])
config_syncer_gcs_to_ndb.SyncInventoryGroupsToNDB()
ndb.get_context().clear_cache()
res = datastore_entities.HostGroupConfig.query(
datastore_entities.HostGroupConfig.lab_name == 'foo').fetch()
for g in res:
if g.name is None:
self.assertIsNone(g)
group_map = {g.name: g for g in res}
self.assertLen(group_map, 10)
self.assertSameElements(group_map['all'].parent_groups, [])
self.assertSameElements(group_map['jump'].parent_groups, ['server'])
self.assertSameElements(group_map['dhcp'].parent_groups, ['server'])
self.assertSameElements(group_map['pxe'].parent_groups, ['server'])
self.assertSameElements(group_map['server'].parent_groups, [])
self.assertSameElements(group_map['dtf'].parent_groups, ['tf'])
self.assertSameElements(group_map['storage_tf'].parent_groups, ['tf'])
def _UpdateHostConfigByInventoryData(lab_name, data):
"""Updates HostConfig inventory_groups.
Args:
lab_name: the lab name.
data: the InventoryData object.
"""
keys = []
entities_from_file = {
host.name: _CreateHostConfigEntityFromHostInventory(lab_name, host)
for host in data.hosts.values()
}
for hostname in entities_from_file:
keys.append(ndb.Key(datastore_entities.HostConfig, hostname))
entities_from_ndb = {}
for entity in ndb.get_multi(keys):
if entity:
entities_from_ndb[entity.hostname] = entity
need_update = {}
for hostname in entities_from_file:
old = entities_from_ndb.get(hostname)
new = entities_from_file[hostname]
if not old:
logging.debug('Creating host config %s', new)
need_update[hostname] = new
elif old.inventory_groups != new.inventory_groups:
old.inventory_groups = new.inventory_groups
logging.debug('Updating host config %s', old)
need_update[hostname] = old
ndb.put_multi(need_update.values())
def _UpdateHostGroupConfigByInventoryData(lab_name, data):
"""Updates HostGroupConfig.
The methods load new HostGroupConfig from inventory file, creates new
HostGroupConfig, updates exist HostGroupConfig and remove obsolete
HostGroupConfig.
Args:
lab_name: the lab name.
data: inventory data which was parsed by the ansible inventory module.
"""
entity_from_file = {
group.key: group for group in [
_CreateHostGroupConfigEntityFromGroupInventory(lab_name, group)
for group in data.groups.values()
]
}
logging.debug('Loaded %d HostGroupConfigs of lab %s',
len(entity_from_file), lab_name)
entity_from_ndb = {
g.key: g for g in datastore_entities.HostGroupConfig.query(
datastore_entities.HostGroupConfig.lab_name == lab_name).fetch()
}
need_update = []
for key, group in entity_from_file.items():
if not _CheckConfigEntitiesEqual(entity_from_ndb.get(key), group):
need_update.append(group)
ndb.put_multi(need_update)
logging.debug('Updated %d HostGroupConfigs.', len(need_update))
need_delete = set(entity_from_ndb.keys()).difference(entity_from_file.keys())
ndb.delete_multi(need_delete)
logging.debug('Removed %d HostGroupConfigs', len(need_delete))
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 _UpdateClusterConfigs(cluster_configs):
"""Update cluster configs in ClusterInfo entity to ndb.
Args:
cluster_configs: a list of cluster configs proto.
"""
logging.debug('Updating cluster configs.')
cluster_config_keys = set()
for cluster_config in cluster_configs:
cluster_config_keys.add(
ndb.Key(datastore_entities.ClusterConfig, cluster_config.cluster_name))
entities = ndb.get_multi(cluster_config_keys)
name_to_cluster = {}
for entity in entities:
if entity:
name_to_cluster[entity.cluster_name] = entity
name_to_entity = {}
for cluster_config in cluster_configs:
new_config_entity = datastore_entities.ClusterConfig.FromMessage(
cluster_config)
if not _CheckConfigEntitiesEqual(
name_to_cluster.get(cluster_config.cluster_name), new_config_entity):
logging.debug('Updating cluster config entity: %s.', new_config_entity)
if cluster_config.cluster_name in name_to_entity:
logging.warning(
'%s has duplicated configs.', cluster_config.cluster_name)
name_to_entity[cluster_config.cluster_name] = new_config_entity
ndb.put_multi(name_to_entity.values())
logging.debug('Cluster configs updated.')
def _SetDeviceRecoveryState(
hostname, device_serial, recovery_state):
"""Set device's recovery state."""
device = GetDevice(hostname, device_serial)
if not device:
logging.error("Device (%s, %s) doesn't exist.", hostname, device_serial)
return
entities_to_update = _BuildDeviceRecoveryState(device, recovery_state)
ndb.put_multi(entities_to_update)
def SyncHarnessImageMetadata():
"""The job to read image manifest from GCR and write to NDB."""
if not env_config.CONFIG.should_sync_harness_image:
return common.HTTP_OK
logging.debug('"should_sync_harness_image" is enabled.')
creds = auth.GetComputeEngineCredentials()
try:
response = requests.get(
url=_LIST_TAGS_URL,
headers={'Authorization': _AUTHORIZATION_TMPL.format(creds.token)})
except requests.exceptions.HTTPError as http_error:
logging.exception(
'Error in http request to get image manifest from GCR.')
raise http_error
response_json = response.json()
manifest = response_json.get('manifest')
if not manifest:
raise KeyError('No valid image manifest is in the response: {}'.format(
response_json))
time_now = datetime.datetime.utcnow()
entities_to_update = []
for digest, data in manifest.items():
current_tags = data.get('tag', [])
analysis_result = _AnalyseTradefedDockerImageTags(current_tags)
historical_tags = []
if analysis_result['is_historical_golden']:
historical_tags.append(_GOLDEN_TAG)
key = ndb.Key(
datastore_entities.TestHarnessImageMetadata,
_IMAGE_METADATA_KEY_TMPL.format(_DEFAULT_TRADEFED_REPO, digest))
time_created = datetime.datetime.utcfromtimestamp(
int(data['timeCreatedMs']) / 1000)
entity = datastore_entities.TestHarnessImageMetadata(
key=key,
repo_name=_DEFAULT_TRADEFED_REPO,
digest=digest,
test_harness=_TRADEFED_HARNESS_NAME,
test_harness_version=analysis_result['tradefed_version_number'],
current_tags=current_tags,
historical_tags=historical_tags,
create_time=time_created,
sync_time=time_now)
entities_to_update.append(entity)
ndb.put_multi(entities_to_update)
# Delete stale image metadata
datastore_util.DeleteEntitiesUpdatedEarlyThanSomeTimeAgo(
datastore_entities.TestHarnessImageMetadata,
datastore_entities.TestHarnessImageMetadata.sync_time,
_STALE_METADATA_MAX_AGE)
return common.HTTP_OK
def _UpdateLabs(clusters):
"""Update lab NDB entities based on hosts.
Args:
clusters: a list of ClusterInfo.
1. Add lab if the lab doesn't exist yet.
2. Refresh the host update state summary in all labs based on the underlying
host groups.
"""
logging.info('Updating labs')
labs_query = datastore_entities.LabInfo.query()
labs_by_lab_names = {lab.lab_name: lab for lab in labs_query}
clusters_by_lab_names = collections.defaultdict(list)
for cluster_info in clusters:
lab_name = cluster_info.lab_name or common.UNKNOWN_LAB_NAME
clusters_by_lab_names[lab_name].append(cluster_info)
labs = []
for lab_name, cluster_infos in clusters_by_lab_names.items():
lab_host_update_state_summary = datastore_entities.HostUpdateStateSummary()
lab_host_update_state_summaries_by_version = collections.defaultdict(
datastore_entities.HostUpdateStateSummary)
host_count_by_harness_version = collections.Counter()
for cluster_info in cluster_infos:
if cluster_info and cluster_info.host_update_state_summary:
lab_host_update_state_summary += cluster_info.host_update_state_summary
if cluster_info and cluster_info.host_update_state_summaries_by_version:
for summary in cluster_info.host_update_state_summaries_by_version:
lab_host_update_state_summaries_by_version[
summary.target_version] += summary
if cluster_info and cluster_info.host_count_by_harness_version:
host_count_by_harness_version += collections.Counter(
cluster_info.host_count_by_harness_version)
if lab_name in labs_by_lab_names:
lab = labs_by_lab_names[lab_name]
else:
lab = datastore_entities.LabInfo(
id=lab_name,
lab_name=lab_name)
lab.populate(
host_update_state_summary=lab_host_update_state_summary,
host_count_by_harness_version=host_count_by_harness_version,
host_update_state_summaries_by_version=list(
lab_host_update_state_summaries_by_version.values()),
update_timestamp=_Now())
labs.append(lab)
ndb.put_multi(labs)
logging.info('Updated labs.')
def testListPredefinedMessages_filtersAndOrdering(self):
"""Test list PredefinedMessages."""
pred_msg_entities = [
datastore_entities.PredefinedMessage(
lab_name='lab-name-1',
type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
content='content-1',
used_count=2),
datastore_entities.PredefinedMessage(
lab_name='lab-name-2',
type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
content='content-2',
used_count=1),
datastore_entities.PredefinedMessage(
lab_name='lab-name-2',
type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
content='content-3',
used_count=3),
datastore_entities.PredefinedMessage(
lab_name='lab-name-4',
type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
content='content-4',
used_count=3),
]
ndb.put_multi(pred_msg_entities)
api_request = {
'type': 'DEVICE_RECOVERY_ACTION',
'lab_name': 'lab-name-2'
}
api_response = self.testapp.post_json(
'/_ah/api/PredefinedMessageApi.ListPredefinedMessages',
api_request)
pred_msgs = protojson.decode_message(
api_messages.PredefinedMessageCollection,
api_response.body).predefined_messages
# The results are filtered by type and lab name.
self.assertEqual(2, len(pred_msgs))
# The results are sorted by count in descending order.
self.assertEqual(pred_msg_entities[1].lab_name, pred_msgs[1].lab_name)
self.assertEqual(pred_msg_entities[1].content, pred_msgs[1].content)
self.assertEqual(pred_msg_entities[1].type, pred_msgs[1].type)
self.assertEqual(pred_msg_entities[1].used_count,
pred_msgs[1].used_count)
self.assertEqual(pred_msg_entities[2].lab_name, pred_msgs[0].lab_name)
self.assertEqual(pred_msg_entities[2].content, pred_msgs[0].content)
self.assertEqual(pred_msg_entities[2].type, pred_msgs[0].type)
self.assertEqual(pred_msg_entities[2].used_count,
pred_msgs[0].used_count)
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
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 _DoUpdateDevicesInNDB(reported_devices, event):
"""Update device entities to ndb.
Args:
reported_devices: device serial to device data mapping.
event: the event have hostname, cluster info and timestamp.
"""
entities_to_update = []
device_keys = []
for device_serial in reported_devices.keys():
device_key = ndb.Key(
datastore_entities.HostInfo, event.hostname,
datastore_entities.DeviceInfo, device_serial)
device_keys.append(device_key)
# If the device doesn't exist, the corresponding entry will be None.
devices = ndb.get_multi(device_keys)
for device, device_key in zip(devices, device_keys):
entities_to_update.extend(
_UpdateDeviceInNDB(
device, device_key, reported_devices.get(device_key.id()),
event))
ndb.put_multi(entities_to_update)
def HideHost(hostname):
"""Hide a host and its devices."""
logging.info("Hide host %s.", hostname)
host = GetHost(hostname)
if not host:
return None
if host.hidden:
logging.info("Host %s is already hidden.", hostname)
return host
now = common.Now()
entities_to_update = []
host.hidden = True
host.timestamp = now
entities_to_update.append(host)
devices = GetDevicesOnHost(hostname)
for device in devices or []:
if device.hidden:
continue
logging.debug("Hide device %s.", device.device_serial)
device.hidden = True
device.timestamp = now
entities_to_update.append(device)
ndb.put_multi(entities_to_update)
return host
def _DoUpdateGoneDevicesInNDB(missing_device_keys, timestamp):
"""Do update gone devices in NDB within transactional."""
entities_to_update = []
devices = ndb.get_multi(missing_device_keys)
for device in devices:
if device.timestamp and device.timestamp > timestamp:
logging.debug("Ignore outdated event.")
continue
if (device.state == common.DeviceState.GONE and
device.timestamp and
device.timestamp >= timestamp - ONE_MONTH):
logging.debug("Ignore gone device.")
continue
if device.timestamp and device.timestamp < timestamp - ONE_MONTH:
device.hidden = True
device.timestamp = timestamp
device_state_history, device_history = _UpdateDeviceState(
device, common.DeviceState.GONE, timestamp)
entities_to_update.append(device)
if device_state_history:
entities_to_update.append(device_state_history)
if device_history:
entities_to_update.append(device_history)
ndb.put_multi(entities_to_update)
def testUpdatePredefinedMessage_failAlreadyExist(self):
predefined_messages = [
datastore_entities.PredefinedMessage(
lab_name='lab1',
type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
content='content-1'),
datastore_entities.PredefinedMessage(
lab_name='lab1',
type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON,
content='content-2'),
]
keys = ndb.put_multi(predefined_messages)
api_request = {
'id': keys[0].id(), # the id of the 1st message
'content': 'content-2', # the content of the 2nd message
}
api_response = self.testapp.post_json(
'/_ah/api/PredefinedMessageApi.UpdatePredefinedMessage',
api_request,
expect_errors=True)
self.assertEqual('409 Conflict', api_response.status)
def AddOrUpdateNote(self, request):
"""Add or update a host note.
Args:
request: an API request.
Returns:
an api_messages.Note.
"""
time_now = datetime.datetime.utcnow()
host_note_entity = datastore_util.GetOrCreateEntity(
datastore_entities.Note,
entity_id=request.id,
hostname=request.hostname,
type=common.NoteType.HOST_NOTE)
host_note_entity.populate(
user=request.user,
message=request.message,
timestamp=time_now,
event_time=request.event_time)
entities_to_update = [host_note_entity]
try:
offline_reason_entity = note_manager.PreparePredefinedMessageForNote(
common.PredefinedMessageType.HOST_OFFLINE_REASON,
message_id=request.offline_reason_id,
lab_name=request.lab_name,
content=request.offline_reason)
except note_manager.InvalidParameterError as err:
raise endpoints.BadRequestException("Invalid offline reason: [%s]" % err)
if offline_reason_entity:
host_note_entity.offline_reason = offline_reason_entity.content
entities_to_update.append(offline_reason_entity)
try:
recovery_action_entity = note_manager.PreparePredefinedMessageForNote(
common.PredefinedMessageType.HOST_RECOVERY_ACTION,
message_id=request.recovery_action_id,
lab_name=request.lab_name,
content=request.recovery_action)
except note_manager.InvalidParameterError as err:
raise endpoints.BadRequestException("Invalid recovery action: [%s]" % err)
if recovery_action_entity:
host_note_entity.recovery_action = recovery_action_entity.content
entities_to_update.append(recovery_action_entity)
keys = ndb.put_multi(entities_to_update)
host_note_msg = datastore_entities.ToMessage(host_note_entity)
host_note_event_msg = api_messages.NoteEvent(
note=host_note_msg, lab_name=request.lab_name)
note_manager.PublishMessage(host_note_event_msg,
common.PublishEventType.HOST_NOTE_EVENT)
note_key = keys[0]
if request.id != note_key.id():
# If ids are different, then a new note is created, we should create
# a history snapshot.
device_manager.CreateAndSaveHostInfoHistoryFromHostNote(
request.hostname, note_key.id())
return host_note_msg
def testListPredefinedMessages_countAndCursor(self):
"""Test list PredefinedMessages."""
pred_msg_entities = [
datastore_entities.PredefinedMessage(
lab_name='lab-name-2',
type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
content='content-1',
used_count=4),
datastore_entities.PredefinedMessage(
lab_name='lab-name-2',
type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
content='content-2',
used_count=3),
datastore_entities.PredefinedMessage(
lab_name='lab-name-2',
type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
content='content-3',
used_count=2),
datastore_entities.PredefinedMessage(
lab_name='lab-name-2',
type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION,
content='content-4',
used_count=1),
]
ndb.put_multi(pred_msg_entities)
# look up the first page
api_request = {
'type': 'DEVICE_RECOVERY_ACTION',
'lab_name': 'lab-name-2',
'count': 2,
}
api_response = self.testapp.post_json(
'/_ah/api/PredefinedMessageApi.ListPredefinedMessages',
api_request)
pred_msg_collection = protojson.decode_message(
api_messages.PredefinedMessageCollection, api_response.body)
pred_msgs = pred_msg_collection.predefined_messages
self.assertEqual(2, len(pred_msgs))
self.assertEqual(pred_msg_entities[0].content, pred_msgs[0].content)
self.assertEqual(pred_msg_entities[1].content, pred_msgs[1].content)
# look up the second page with next_cursor
api_request = {
'type': 'DEVICE_RECOVERY_ACTION',
'lab_name': 'lab-name-2',
'count': 2,
'cursor': pred_msg_collection.next_cursor,
}
api_response = self.testapp.post_json(
'/_ah/api/PredefinedMessageApi.ListPredefinedMessages',
api_request)
pred_msg_collection = protojson.decode_message(
api_messages.PredefinedMessageCollection, api_response.body)
pred_msgs = pred_msg_collection.predefined_messages
self.assertEqual(2, len(pred_msgs))
self.assertEqual(pred_msg_entities[2].content, pred_msgs[0].content)
self.assertEqual(pred_msg_entities[3].content, pred_msgs[1].content)
# look up the first page again with prev_cursor of the second page
api_request = {
'type': 'DEVICE_RECOVERY_ACTION',
'lab_name': 'lab-name-2',
'count': 2,
'cursor': pred_msg_collection.prev_cursor,
'backwards': True,
}
api_response = self.testapp.post_json(
'/_ah/api/PredefinedMessageApi.ListPredefinedMessages',
api_request)
pred_msg_collection = protojson.decode_message(
api_messages.PredefinedMessageCollection, api_response.body)
pred_msgs = pred_msg_collection.predefined_messages
self.assertEqual(2, len(pred_msgs))
self.assertEqual(pred_msg_entities[0].content, pred_msgs[0].content)
self.assertEqual(pred_msg_entities[1].content, pred_msgs[1].content)
def _UpdateClusters(hosts):
"""Update cluster NDB entities based on hosts.
Args:
hosts: list of HostInfo entity with required field, from the entire system.
Returns:
list of ClusterInfo, clusters to upsert.
"""
logging.info('Updating clusters')
cluster_to_hosts = collections.defaultdict(list)
for host in hosts:
cluster_to_hosts[host.physical_cluster].append(host)
clusters_to_delete = []
clusters_to_upsert = []
query = datastore_entities.ClusterInfo.query()
for cluster in query:
if cluster.cluster not in cluster_to_hosts:
clusters_to_delete.append(cluster.key)
ndb.delete_multi(clusters_to_delete)
logging.debug('Deleted clusters due to no hosts: %s', clusters_to_delete)
query = datastore_entities.HostUpdateState.query()
update_states_by_hostname = {
update_state.hostname: update_state for update_state in query.fetch()}
for cluster, hosts in six.iteritems(cluster_to_hosts):
cluster_id = cluster or common.UNKNOWN_CLUSTER_NAME
cluster_entity = datastore_entities.ClusterInfo(id=cluster_id)
for host in hosts:
if host.lab_name:
cluster_entity.lab_name = host.lab_name
break
else:
cluster_entity.lab_name = common.UNKNOWN_LAB_NAME
cluster_entity.cluster = cluster_id
cluster_entity.total_devices = 0
cluster_entity.offline_devices = 0
cluster_entity.available_devices = 0
cluster_entity.allocated_devices = 0
cluster_entity.device_count_timestamp = _Now()
host_update_states = []
host_update_states_by_target_version = collections.defaultdict(list)
host_count_by_harness_version = collections.Counter()
for host in hosts:
cluster_entity.total_devices += host.total_devices or 0
cluster_entity.offline_devices += host.offline_devices or 0
cluster_entity.available_devices += host.available_devices or 0
cluster_entity.allocated_devices += host.allocated_devices or 0
host_update_state = update_states_by_hostname.get(host.hostname)
if host_update_state:
host_update_states.append(host_update_state)
host_update_states_by_target_version[
host_update_state.target_version].append(host_update_state)
if host.test_harness_version:
host_count_by_harness_version[host.test_harness_version] += 1
else:
host_count_by_harness_version[common.UNKNOWN_TEST_HARNESS_VERSION] += 1
cluster_entity.host_count_by_harness_version = host_count_by_harness_version
cluster_entity.host_update_state_summary = _CreateHostUpdateStateSummary(
host_update_states)
for version, states in host_update_states_by_target_version.items():
cluster_entity.host_update_state_summaries_by_version.append(
_CreateHostUpdateStateSummary(states, target_version=version))
clusters_to_upsert.append(cluster_entity)
ndb.put_multi(clusters_to_upsert)
logging.debug('Updated clusters.')
return clusters_to_upsert
def _MarkHostUpdateStateIfTimedOut(hostname):
"""Mark HostUpdateState as TIMED_OUT if it times out.
Args:
hostname: text, the host to check the update timeouts.
Returns:
An instance of HostUpdateState entity, None if it does not exist previously.
"""
host_update_state = datastore_entities.HostUpdateState.get_by_id(hostname)
if not host_update_state:
logging.info('No update state is found for host: %s.', hostname)
return
timeout_timedelta = _GetCustomizedOrDefaultHostUpdateTimeout(hostname)
now = _Now()
entities_to_update = []
if (host_update_state.state and
host_update_state.state in common.NON_FINAL_HOST_UPDATE_STATES):
if host_update_state.update_timestamp:
update_state_age = now - host_update_state.update_timestamp
if timeout_timedelta < update_state_age:
display_message = (
_TIMEDOUT_DISPLAY_MESSAGE_TMPL % (
hostname, host_update_state.state,
host_update_state.update_timestamp,
update_state_age.total_seconds(),
timeout_timedelta.total_seconds()))
logging.info('Marking update state as TIMED_OUT: %s', display_message)
host_update_state.state = api_messages.HostUpdateState.TIMED_OUT
host_update_state.update_timestamp = now
host_update_state.populate(
state=api_messages.HostUpdateState.TIMED_OUT,
update_timestamp=now,
display_message=display_message)
entities_to_update.append(host_update_state)
host_update_state_history = datastore_entities.HostUpdateStateHistory(
parent=ndb.Key(datastore_entities.HostInfo, hostname),
hostname=host_update_state.hostname,
state=host_update_state.state,
update_timestamp=now,
update_task_id=host_update_state.update_task_id,
display_message=display_message)
entities_to_update.append(host_update_state_history)
else:
logging.debug('Host<%s> is in HostUpdateState<%s> since %s.',
hostname, host_update_state.state,
host_update_state.update_timestamp)
else:
logging.debug('Host<%s> has no timestamp in the HostUpdateState. '
'Auto adding a timestamp on it.',
hostname)
host_update_state.update_timestamp = now
entities_to_update.append(host_update_state)
host_update_state_history = datastore_entities.HostUpdateStateHistory(
parent=ndb.Key(datastore_entities.HostInfo, hostname),
hostname=host_update_state.hostname,
state=host_update_state.state,
update_timestamp=now,
update_task_id=host_update_state.update_task_id)
entities_to_update.append(host_update_state_history)
ndb.put_multi(entities_to_update)
return host_update_state
def BatchUpdateNotesWithPredefinedMessage(self, request):
"""Batch update notes with the same predefined message.
Args:
request: an API request.
Returns:
an api_messages.NoteCollection object.
"""
time_now = datetime.datetime.utcnow()
host_note_entities = []
for note in request.notes:
note_id = int(note.id) if note.id is not None else None
host_note_entity = datastore_util.GetOrCreateEntity(
datastore_entities.Note,
entity_id=note_id,
hostname=note.hostname,
type=common.NoteType.HOST_NOTE)
host_note_entity.populate(
user=request.user,
message=request.message,
timestamp=time_now,
event_time=request.event_time)
host_note_entities.append(host_note_entity)
try:
offline_reason_entity = note_manager.PreparePredefinedMessageForNote(
common.PredefinedMessageType.HOST_OFFLINE_REASON,
message_id=request.offline_reason_id,
lab_name=request.lab_name,
content=request.offline_reason,
delta_count=len(host_note_entities))
except note_manager.InvalidParameterError as err:
raise endpoints.BadRequestException("Invalid offline reason: [%s]" % err)
if offline_reason_entity:
for host_note_entity in host_note_entities:
host_note_entity.offline_reason = offline_reason_entity.content
offline_reason_entity.put()
try:
recovery_action_entity = note_manager.PreparePredefinedMessageForNote(
common.PredefinedMessageType.HOST_RECOVERY_ACTION,
message_id=request.recovery_action_id,
lab_name=request.lab_name,
content=request.recovery_action,
delta_count=len(host_note_entities))
except note_manager.InvalidParameterError as err:
raise endpoints.BadRequestException("Invalid recovery action: [%s]" % err)
if recovery_action_entity:
for host_note_entity in host_note_entities:
host_note_entity.recovery_action = recovery_action_entity.content
recovery_action_entity.put()
note_keys = ndb.put_multi(host_note_entities)
host_note_entities = ndb.get_multi(note_keys)
note_msgs = []
for host_note_entity in host_note_entities:
host_note_msg = datastore_entities.ToMessage(host_note_entity)
note_msgs.append(host_note_msg)
host_note_event_msg = api_messages.NoteEvent(
note=host_note_msg,
lab_name=request.lab_name)
note_manager.PublishMessage(
host_note_event_msg, common.PublishEventType.HOST_NOTE_EVENT)
for request_note, updated_note_key in zip(request.notes, note_keys):
if not request_note.id:
# If ids are not provided, then a new note is created, we should create
# a history snapshot.
device_manager.CreateAndSaveHostInfoHistoryFromHostNote(
request_note.hostname, updated_note_key.id())
return api_messages.NoteCollection(
notes=note_msgs, more=False, next_cursor=None, prev_cursor=None)
def BatchUpdateHostMetadata(self, request):
"""Update HostMetadata on multiple hosts.
Args:
request: an API request.
Request Params:
hostname: list of strings, the name of hosts.
test_harness_image: string, the url to test harness image.
user: string, the user sending the request.
Returns:
a message_types.VoidMessage object.
Raises:
endpoints.BadRequestException, when request does not match existing hosts.
"""
host_configs = ndb.get_multi(
ndb.Key(datastore_entities.HostConfig, hostname)
for hostname in request.hostnames)
host_metadatas = ndb.get_multi(
ndb.Key(datastore_entities.HostMetadata, hostname)
for hostname in request.hostnames)
hosts_no_permission = []
hosts_not_enabled = []
metadatas_to_update = []
for hostname, config, metadata in zip(
request.hostnames, host_configs, host_metadatas):
if not config or not config.enable_ui_update:
hosts_not_enabled.append(hostname)
continue
if request.user not in config.owners:
hosts_no_permission.append(hostname)
continue
if not metadata:
metadata = datastore_entities.HostMetadata(
id=hostname, hostname=hostname)
if not harness_image_metadata_syncer.AreHarnessImagesEqual(
metadata.test_harness_image, request.test_harness_image):
event = host_event.HostEvent(
time=datetime.datetime.utcnow(),
type=_HOST_UPDATE_STATE_CHANGED_EVENT_NAME,
hostname=hostname,
host_update_state=_HOST_UPDATE_STATE_PENDING,
data={"host_update_target_image": request.test_harness_image})
device_manager.HandleDeviceSnapshotWithNDB(event)
metadata.populate(test_harness_image=request.test_harness_image)
metadatas_to_update.append(metadata)
ndb.put_multi(metadatas_to_update)
if not hosts_no_permission and not hosts_not_enabled:
return message_types.VoidMessage()
error_message = ""
if hosts_no_permission:
error_message += (
"Request user %s is not in the owner list of hosts [%s]. "
% (request.user, ", ".join(hosts_no_permission)))
if hosts_not_enabled:
error_message += ("Hosts [%s] are not enabled to be updated from UI. "
% ", ".join(hosts_not_enabled))
raise endpoints.BadRequestException(error_message)
def testSyncHarnessImageMetadata_OverwriteExistingEntities(
self, mock_requests, mock_util_datetime, mock_syncer_datetime,
mock_auth):
"""Test sync harness image metadata."""
time_now = datetime.datetime(2020, 12, 24)
time_created = datetime.datetime(2020, 12, 10)
time_created_ms = str(
int((time_created - datetime.datetime(1970, 1, 1)).total_seconds()
* 1000))
mock_util_datetime.datetime.utcnow.return_value = time_now
mock_syncer_datetime.datetime.utcnow.return_value = time_now
mock_syncer_datetime.datetime.utcfromtimestamp = (
datetime.datetime.utcfromtimestamp)
mock_requests.get().json.return_value = {
'manifest': {
'sha1': {
'tag': [
'111111',
'golden',
'canary',
'golden_tradefed_image_20201210_1200_RC00',
],
'timeCreatedMs':
time_created_ms,
},
'sha2': {
'tag': [
'2222222',
'golden_tradefed_image_20201210_0600_RC00',
],
'timeCreatedMs': time_created_ms,
},
'sha3': {
'tag': [
'3333333',
'staging',
],
'timeCreatedMs': time_created_ms,
},
}
}
existing_entities = [
datastore_entities.TestHarnessImageMetadata(
key=ndb.Key(datastore_entities.TestHarnessImageMetadata,
'gcr.io/dockerized-tradefed/tradefed:sha1'),
repo_name='gcr.io/dockerized-tradefed/tradefed',
digest='sha1',
test_harness=harness_image_metadata_syncer.
_TRADEFED_HARNESS_NAME,
test_harness_version='111111',
current_tags=['111111', 'canary', 'staging'],
create_time=time_created,
sync_time=time_now),
datastore_entities.TestHarnessImageMetadata(
key=ndb.Key(datastore_entities.TestHarnessImageMetadata,
'gcr.io/dockerized-tradefed/tradefed:sha2'),
repo_name='gcr.io/dockerized-tradefed/tradefed',
digest='sha2',
test_harness=harness_image_metadata_syncer.
_TRADEFED_HARNESS_NAME,
test_harness_version='2222222',
current_tags=[
'2222222', 'golden',
'golden_tradefed_image_20201210_0600_RC00'
],
historical_tags=['golden'],
create_time=time_created,
sync_time=time_now),
]
ndb.put_multi(existing_entities)
harness_image_metadata_syncer.SyncHarnessImageMetadata()
keys = [
ndb.Key(datastore_entities.TestHarnessImageMetadata,
'gcr.io/dockerized-tradefed/tradefed:sha1'),
ndb.Key(datastore_entities.TestHarnessImageMetadata,
'gcr.io/dockerized-tradefed/tradefed:sha2'),
ndb.Key(datastore_entities.TestHarnessImageMetadata,
'gcr.io/dockerized-tradefed/tradefed:sha3'),
]
entity_1, entity_2, entity_3 = ndb.get_multi(keys)
self.assertEqual('sha1', entity_1.digest)
self.assertEqual('111111', entity_1.test_harness_version)
self.assertEqual(time_created, entity_1.create_time)
self.assertEqual(time_now, entity_1.sync_time)
self.assertCountEqual([
'111111', 'golden', 'canary',
'golden_tradefed_image_20201210_1200_RC00'
], entity_1.current_tags)
self.assertCountEqual(['golden'], entity_1.historical_tags)
self.assertEqual('sha2', entity_2.digest)
self.assertEqual('2222222', entity_2.test_harness_version)
self.assertEqual(time_created, entity_2.create_time)
self.assertEqual(time_now, entity_2.sync_time)
self.assertCountEqual([
'2222222',
'golden_tradefed_image_20201210_0600_RC00',
], entity_2.current_tags)
self.assertCountEqual(['golden'], entity_2.historical_tags)
self.assertEqual('sha3', entity_3.digest)
self.assertEqual('3333333', entity_3.test_harness_version)
self.assertEqual(time_created, entity_3.create_time)
self.assertEqual(time_now, entity_3.sync_time)
self.assertCountEqual(['3333333', 'staging'], entity_3.current_tags)
self.assertEmpty(entity_3.historical_tags)
