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