autotest: Make 'get_hosts' respect static attributes & labels.
BUG=chromium:792309
TEST=Ran unittest. Call RPC 'get_hosts' locally.
Change-Id: I1fd5cbe59781fd3c17816c5e386cffc05635c0ff
Reviewed-on: https://chromium-review.googlesource.com/929609
Commit-Ready: Xixuan Wu <[email protected]>
Tested-by: Xixuan Wu <[email protected]>
Reviewed-by: Shuqian Zhao <[email protected]>
diff --git a/frontend/afe/rpc_interface.py b/frontend/afe/rpc_interface.py
index d086d5d..a8b4302 100644
--- a/frontend/afe/rpc_interface.py
+++ b/frontend/afe/rpc_interface.py
@@ -89,6 +89,9 @@
RESPECT_STATIC_LABELS = global_config.global_config.get_config_value(
'SKYLAB', 'respect_static_labels', type=bool, default=False)
+RESPECT_STATIC_ATTRIBUTES = global_config.global_config.get_config_value(
+ 'SKYLAB', 'respect_static_attributes', type=bool, default=False)
+
# Relevant CrosDynamicSuiteExceptions are defined in client/common_lib/error.py.
# labels
@@ -637,14 +640,40 @@
'acl_list')
models.Host.objects.populate_relationships(hosts, models.HostAttribute,
'attribute_list')
+ models.Host.objects.populate_relationships(hosts,
+ models.StaticHostAttribute,
+ 'staticattribute_list')
host_dicts = []
for host_obj in hosts:
host_dict = host_obj.get_object_dict()
- host_dict['labels'] = [label.name for label in host_obj.label_list]
- host_dict['platform'] = rpc_utils.find_platform(host_obj)
host_dict['acls'] = [acl.name for acl in host_obj.acl_list]
host_dict['attributes'] = dict((attribute.attribute, attribute.value)
for attribute in host_obj.attribute_list)
+ if RESPECT_STATIC_LABELS:
+ label_list = []
+ # Only keep static labels which has a corresponding entries in
+ # afe_labels.
+ for label in host_obj.label_list:
+ if label.is_replaced_by_static():
+ static_label = models.StaticLabel.smart_get(label.name)
+ label_list.append(static_label)
+ else:
+ label_list.append(label)
+
+ host_dict['labels'] = [label.name for label in label_list]
+ host_dict['platform'] = rpc_utils.find_platform(
+ host_obj.hostname, label_list)
+ else:
+ host_dict['labels'] = [label.name for label in host_obj.label_list]
+ host_dict['platform'] = rpc_utils.find_platform(
+ host_obj.hostname, host_obj.label_list)
+
+ if RESPECT_STATIC_ATTRIBUTES:
+ # Overwrite attribute with values in afe_static_host_attributes.
+ for attr in host_obj.staticattribute_list:
+ if attr.attribute in host_dict['attributes']:
+ host_dict['attributes'][attr.attribute] = attr.value
+
if include_current_job:
host_dict['current_job'] = None
host_dict['current_special_task'] = None
@@ -660,6 +689,7 @@
'%d-%s' % (tasks[0].get_object_dict()['id'],
tasks[0].get_object_dict()['task'].lower()))
host_dicts.append(host_dict)
+
return rpc_utils.prepare_for_serialization(host_dicts)
diff --git a/frontend/afe/rpc_interface_unittest.py b/frontend/afe/rpc_interface_unittest.py
index f134d97..017b122 100755
--- a/frontend/afe/rpc_interface_unittest.py
+++ b/frontend/afe/rpc_interface_unittest.py
@@ -246,6 +246,68 @@
self._do_heartbeat_and_assert_response(known_hosts=[host1])
+class RpcInterfaceTestWithStaticAttribute(
+ mox.MoxTestBase, unittest.TestCase,
+ frontend_test_utils.FrontendTestMixin):
+
+ def setUp(self):
+ super(RpcInterfaceTestWithStaticAttribute, self).setUp()
+ self._frontend_common_setup()
+ self.god = mock.mock_god()
+ self.old_respect_static_config = rpc_interface.RESPECT_STATIC_ATTRIBUTES
+ rpc_interface.RESPECT_STATIC_ATTRIBUTES = True
+ models.RESPECT_STATIC_ATTRIBUTES = True
+
+
+ def tearDown(self):
+ self.god.unstub_all()
+ self._frontend_common_teardown()
+ global_config.global_config.reset_config_values()
+ rpc_interface.RESPECT_STATIC_ATTRIBUTES = self.old_respect_static_config
+ models.RESPECT_STATIC_ATTRIBUTES = self.old_respect_static_config
+
+
+ def _set_static_attribute(self, host, attribute, value):
+ """Set static attribute for a host.
+
+ It ensures that all static attributes have a corresponding
+ entry in afe_host_attributes.
+ """
+ # Get or create the reference object in afe_host_attributes.
+ model, args = host._get_attribute_model_and_args(attribute)
+ model.objects.get_or_create(**args)
+
+ attribute_model, get_args = host._get_static_attribute_model_and_args(
+ attribute)
+ attribute_object, _ = attribute_model.objects.get_or_create(**get_args)
+ attribute_object.value = value
+ attribute_object.save()
+
+
+ def _fake_host_with_static_attributes(self):
+ host1 = models.Host.objects.create(hostname='test_host')
+ host1.set_attribute('test_attribute1', 'test_value1')
+ host1.set_attribute('test_attribute2', 'test_value2')
+ self._set_static_attribute(host1, 'test_attribute1', 'static_value1')
+ self._set_static_attribute(host1, 'static_attribute1', 'static_value2')
+ host1.save()
+ return host1
+
+
+ def test_get_hosts(self):
+ host1 = self._fake_host_with_static_attributes()
+ hosts = rpc_interface.get_hosts(hostname=host1.hostname)
+ host = hosts[0]
+
+ self.assertEquals(host['hostname'], 'test_host')
+ self.assertEquals(host['acls'], ['Everyone'])
+ # Respect the value of static attributes.
+ self.assertEquals(host['attributes'],
+ {'test_attribute1': 'static_value1',
+ 'test_attribute2': 'test_value2',
+ 'static_attribute1': 'static_value2'})
+
+
class RpcInterfaceTestWithStaticLabel(ShardHeartbeatTest,
frontend_test_utils.FrontendTestMixin):
@@ -268,6 +330,44 @@
models.RESPECT_STATIC_LABELS = self.old_respect_static_config
+ def _fake_host_with_static_labels(self):
+ host1 = models.Host.objects.create(hostname='test_host')
+ label1 = models.Label.objects.create(
+ name='non_static_label1', platform=False)
+ non_static_platform = models.Label.objects.create(
+ name='static_platform', platform=False)
+ static_platform = models.StaticLabel.objects.create(
+ name='static_platform', platform=True)
+ models.ReplacedLabel.objects.create(label_id=non_static_platform.id)
+ host1.static_labels.add(static_platform)
+ host1.labels.add(non_static_platform)
+ host1.labels.add(label1)
+ host1.save()
+ return host1
+
+
+ def test_get_hosts(self):
+ host1 = self._fake_host_with_static_labels()
+ hosts = rpc_interface.get_hosts(hostname=host1.hostname)
+ host = hosts[0]
+
+ self.assertEquals(host['hostname'], 'test_host')
+ self.assertEquals(host['acls'], ['Everyone'])
+ # Respect all labels in afe_hosts_labels.
+ self.assertEquals(host['labels'],
+ ['non_static_label1', 'static_platform'])
+ # Respect static labels.
+ self.assertEquals(host['platform'], 'static_platform')
+
+
+ def test_get_hosts_multiple_labels(self):
+ self._fake_host_with_static_labels()
+ hosts = rpc_interface.get_hosts(
+ multiple_labels=['non_static_label1', 'static_platform'])
+ host = hosts[0]
+ self.assertEquals(host['hostname'], 'test_host')
+
+
def test_delete_static_label(self):
label1 = models.Label.smart_get('static')
diff --git a/frontend/afe/rpc_utils.py b/frontend/afe/rpc_utils.py
index 5741d2e..c347ffb 100644
--- a/frontend/afe/rpc_utils.py
+++ b/frontend/afe/rpc_utils.py
@@ -557,21 +557,26 @@
return False
-def find_platform(host):
+def find_platform(hostname, label_list):
"""
Figure out the platform name for the given host
object. If none, the return value for either will be None.
+ @param hostname: The hostname to find platform.
+ @param label_list: The label list to find platform.
+
@returns platform name for the given host.
"""
- platforms = [label.name for label in host.label_list if label.platform]
+ platforms = [label.name for label in label_list if label.platform]
if not platforms:
platform = None
else:
platform = platforms[0]
+
if len(platforms) > 1:
raise ValueError('Host %s has more than one platform: %s' %
- (host.hostname, ', '.join(platforms)))
+ (hostname, ', '.join(platforms)))
+
return platform
