diff --git a/nova/compute/provider_tree.py b/nova/compute/provider_tree.py index d688577789..2324228291 100644 --- a/nova/compute/provider_tree.py +++ b/nova/compute/provider_tree.py @@ -533,6 +533,30 @@ class ProviderTree(object): provider = self._find_with_lock(name_or_uuid) return provider.update_traits(traits, generation=generation) + def add_traits(self, name_or_uuid, *traits): + """Set traits on a provider, without affecting existing traits. + + :param name_or_uuid: The name or UUID of the provider whose traits are + to be affected. + :param traits: String names of traits to be added. + """ + with self.lock: + provider = self._find_with_lock(name_or_uuid) + final_traits = provider.traits | set(traits) + provider.update_traits(final_traits) + + def remove_traits(self, name_or_uuid, *traits): + """Unset traits on a provider, without affecting other existing traits. + + :param name_or_uuid: The name or UUID of the provider whose traits are + to be affected. + :param traits: String names of traits to be removed. + """ + with self.lock: + provider = self._find_with_lock(name_or_uuid) + final_traits = provider.traits - set(traits) + provider.update_traits(final_traits) + def in_aggregates(self, name_or_uuid, aggregates): """Given a name or UUID of a provider, query whether that provider is a member of *all* the specified aggregates. @@ -588,3 +612,28 @@ class ProviderTree(object): provider = self._find_with_lock(name_or_uuid) return provider.update_aggregates(aggregates, generation=generation) + + def add_aggregates(self, name_or_uuid, *aggregates): + """Set aggregates on a provider, without affecting existing aggregates. + + :param name_or_uuid: The name or UUID of the provider whose aggregates + are to be affected. + :param aggregates: String UUIDs of aggregates to be added. + """ + with self.lock: + provider = self._find_with_lock(name_or_uuid) + final_aggs = provider.aggregates | set(aggregates) + provider.update_aggregates(final_aggs) + + def remove_aggregates(self, name_or_uuid, *aggregates): + """Unset aggregates on a provider, without affecting other existing + aggregates. + + :param name_or_uuid: The name or UUID of the provider whose aggregates + are to be affected. + :param aggregates: String UUIDs of aggregates to be removed. + """ + with self.lock: + provider = self._find_with_lock(name_or_uuid) + final_aggs = provider.aggregates - set(aggregates) + provider.update_aggregates(final_aggs) diff --git a/nova/tests/unit/compute/test_provider_tree.py b/nova/tests/unit/compute/test_provider_tree.py index cc0b8a238b..115e37561d 100644 --- a/nova/tests/unit/compute/test_provider_tree.py +++ b/nova/tests/unit/compute/test_provider_tree.py @@ -552,6 +552,36 @@ class TestProviderTree(test.NoDBTestCase): self.assertEqual(rp_gen, pt.data(cn.uuid).generation) self.assertTrue(pt.has_traits(cn.uuid, traits[-1:])) + def test_add_remove_traits(self): + cn = self.compute_node1 + pt = self._pt_with_cns() + self.assertEqual(set([]), pt.data(cn.uuid).traits) + # Add a couple of traits + pt.add_traits(cn.uuid, "HW_GPU_API_DIRECT3D_V7_0", "HW_NIC_OFFLOAD_SG") + self.assertEqual( + set(["HW_GPU_API_DIRECT3D_V7_0", "HW_NIC_OFFLOAD_SG"]), + pt.data(cn.uuid).traits) + # set() behavior: add a trait that's already there, and one that's not. + # The unrelated one is unaffected. + pt.add_traits(cn.uuid, "HW_GPU_API_DIRECT3D_V7_0", "HW_CPU_X86_AVX") + self.assertEqual( + set(["HW_GPU_API_DIRECT3D_V7_0", "HW_NIC_OFFLOAD_SG", + "HW_CPU_X86_AVX"]), + pt.data(cn.uuid).traits) + # Now remove a trait + pt.remove_traits(cn.uuid, "HW_NIC_OFFLOAD_SG") + self.assertEqual( + set(["HW_GPU_API_DIRECT3D_V7_0", "HW_CPU_X86_AVX"]), + pt.data(cn.uuid).traits) + # set() behavior: remove a trait that's there, and one that's not. + # The unrelated one is unaffected. + pt.remove_traits(cn.uuid, + "HW_NIC_OFFLOAD_SG", "HW_GPU_API_DIRECT3D_V7_0") + self.assertEqual(set(["HW_CPU_X86_AVX"]), pt.data(cn.uuid).traits) + # Remove the last trait, and an unrelated one + pt.remove_traits(cn.uuid, "CUSTOM_FOO", "HW_CPU_X86_AVX") + self.assertEqual(set([]), pt.data(cn.uuid).traits) + def test_have_aggregates_changed_no_existing_rp(self): pt = self._pt_with_cns() self.assertRaises( @@ -606,3 +636,29 @@ class TestProviderTree(test.NoDBTestCase): self.assertTrue(pt.in_aggregates(cn.uuid, aggregates[-1:])) # Previously-taken data now differs self.assertTrue(pt.have_aggregates_changed(cn.uuid, cnsnap.aggregates)) + + def test_add_remove_aggregates(self): + cn = self.compute_node1 + pt = self._pt_with_cns() + self.assertEqual(set([]), pt.data(cn.uuid).aggregates) + # Add a couple of aggregates + pt.add_aggregates(cn.uuid, uuids.agg1, uuids.agg2) + self.assertEqual( + set([uuids.agg1, uuids.agg2]), + pt.data(cn.uuid).aggregates) + # set() behavior: add an aggregate that's already there, and one that's + # not. The unrelated one is unaffected. + pt.add_aggregates(cn.uuid, uuids.agg1, uuids.agg3) + self.assertEqual(set([uuids.agg1, uuids.agg2, uuids.agg3]), + pt.data(cn.uuid).aggregates) + # Now remove an aggregate + pt.remove_aggregates(cn.uuid, uuids.agg2) + self.assertEqual(set([uuids.agg1, uuids.agg3]), + pt.data(cn.uuid).aggregates) + # set() behavior: remove an aggregate that's there, and one that's not. + # The unrelated one is unaffected. + pt.remove_aggregates(cn.uuid, uuids.agg2, uuids.agg3) + self.assertEqual(set([uuids.agg1]), pt.data(cn.uuid).aggregates) + # Remove the last aggregate, and an unrelated one + pt.remove_aggregates(cn.uuid, uuids.agg4, uuids.agg1) + self.assertEqual(set([]), pt.data(cn.uuid).aggregates)