PBM related utility methods

This change adds following PBM related utility methods:
* Get all profiles
* Get profile ID by profile name
* Filter hubs by profile
* Convert datastores to hubs
* Filter datastores by hubs

Change-Id: I47524f28edfa4325340046cf266184771ea4e3dc

2 files changed, 239 insertions, 1 deletions

diff --git a/oslo/vmware/pbm.py b/oslo/vmware/pbm.py
index 004c881..dbd65cc 100644
--- a/oslo/vmware/pbm.py
+++ b/oslo/vmware/pbm.py
@@ -14,15 +14,17 @@
 #    under the License.
 
 """
-VMware PBM client.
+VMware PBM client and PBM related utility methods
 
 PBM is used for policy based placement in VMware datastores.
 Refer http://goo.gl/GR2o6U for more details.
 """
 
+import logging
 import suds
 import suds.sax.element as element
 
+from oslo.vmware.openstack.common.gettextutils import _
 from oslo.vmware import vim
 from oslo.vmware import vim_util
 
@@ -30,6 +32,8 @@ from oslo.vmware import vim_util
 SERVICE_INSTANCE = 'ServiceInstance'
 SERVICE_TYPE = 'PbmServiceInstance'
 
+LOG = logging.getLogger(__name__)
+
 
 class PBMClient(vim.Vim):
     """SOAP based PBM client."""
@@ -64,3 +68,102 @@ class PBMClient(vim.Vim):
             self._pbm_service_content = (
                 self._pbm_client.service.PbmRetrieveServiceContent(si_moref))
         return self._pbm_service_content
+
+
+def get_all_profiles(session):
+    """Get all the profiles defined in VC server.
+
+    :returns: PbmProfile data objects
+    :raises: VimException, VimFaultException, VimAttributeException,
+             VimSessionOverLoadException, VimConnectionException
+    """
+    LOG.debug(_("Fetching all the profiles defined in VC server."))
+
+    pbm = session.pbm
+    profile_manager = pbm.service_content.profileManager
+    res_type = pbm.client.factory.create('ns0:PbmProfileResourceType')
+    res_type.resourceType = 'STORAGE'
+
+    profile_ids = session.invoke_api(pbm,
+                                     'PbmQueryProfile',
+                                     profile_manager,
+                                     resourceType=res_type)
+    LOG.debug(_("Fetched profile IDs: %s."), profile_ids)
+    return session.invoke_api(pbm,
+                              'PbmRetrieveContent',
+                              profile_manager,
+                              profileIds=profile_ids)
+
+
+def get_profile_id_by_name(session, profile_name):
+    """Get the profile UUID corresponding to the given profile name.
+
+    :param profile_name: profile name whose UUID needs to be retrieved
+    :returns: profile UUID string or None if profile not found
+    :raises: VimException, VimFaultException, VimAttributeException,
+             VimSessionOverLoadException, VimConnectionException
+    """
+    LOG.debug(_("Retrieving profile ID for profile: %s."), profile_name)
+    for profile in get_all_profiles(session):
+        if profile.name == profile_name:
+            profile_id = profile.profileId
+            LOG.debug(_("Retrieved profile ID: %(id)s for profile: %(name)s."),
+                      {'id': profile_id,
+                       'name': profile_name})
+            return profile_id
+    return None
+
+
+def filter_hubs_by_profile(session, hubs, profile_id):
+    """Filter and return hubs that match the given profile.
+
+    :param hubs: PbmPlacementHub morefs
+    :param profile_id: profile ID
+    :returns: subset of hubs that match the given profile
+    :raises: VimException, VimFaultException, VimAttributeException,
+             VimSessionOverLoadException, VimConnectionException
+    """
+    LOG.debug(_("Filtering hubs: %(hubs)s that match profile: %(profile)s."),
+              {'hubs': hubs,
+               'profile': profile_id})
+
+    pbm = session.pbm
+    placement_solver = pbm.service_content.placementSolver
+    filtered_hubs = session.invoke_api(pbm,
+                                       'PbmQueryMatchingHub',
+                                       placement_solver,
+                                       hubsToSearch=hubs,
+                                       profile=profile_id)
+    LOG.debug(_("Filtered hubs: %s"), filtered_hubs)
+    return filtered_hubs
+
+
+def convert_datastores_to_hubs(pbm_client_factory, datastores):
+    """Convert given datastore morefs to PbmPlacementHub morefs.
+
+    :param pbm_client_factory: Factory to create PBM API input specs
+    :param datastores: list of datastore morefs
+    :returns: list of PbmPlacementHub morefs
+    """
+    hubs = []
+    for ds in datastores:
+        hub = pbm_client_factory.create('ns0:PbmPlacementHub')
+        hub.hubId = ds.value
+        hub.hubType = 'Datastore'
+        hubs.append(hub)
+    return hubs
+
+
+def filter_datastores_by_hubs(hubs, datastores):
+    """Get filtered subset of datastores corresponding to the given hub list.
+
+    :param hubs: list of PbmPlacementHub morefs
+    :param datastores: all candidate datastores
+    :returns: subset of datastores corresponding to the given hub list
+    """
+    filtered_dss = []
+    hub_ids = [hub.hubId for hub in hubs]
+    for ds in datastores:
+        if ds.value in hub_ids:
+            filtered_dss.append(ds)
+    return filtered_dss
diff --git a/tests/test_pbm.py b/tests/test_pbm.py
new file mode 100644
index 0000000..a505443
--- /dev/null
+++ b/tests/test_pbm.py
@@ -0,0 +1,135 @@
+# Copyright (c) 2014 VMware, Inc.
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""
+Unit tests for PBM utility methods.
+"""
+
+import mock
+
+from oslo.vmware import pbm
+from tests import base
+
+
+class PBMUtilityTest(base.TestCase):
+    """Tests for PBM utility methods."""
+
+    def test_get_all_profiles(self):
+        session = mock.Mock()
+        session.pbm = mock.Mock()
+        profile_ids = mock.Mock()
+
+        def invoke_api_side_effect(module, method, *args, **kwargs):
+            self.assertEqual(session.pbm, module)
+            self.assertTrue(method in ['PbmQueryProfile',
+                                       'PbmRetrieveContent'])
+            self.assertEqual(session.pbm.service_content.profileManager,
+                             args[0])
+            if method == 'PbmQueryProfile':
+                self.assertEqual('STORAGE',
+                                 kwargs['resourceType'].resourceType)
+                return profile_ids
+            self.assertEqual(profile_ids, kwargs['profileIds'])
+
+        session.invoke_api.side_effect = invoke_api_side_effect
+        pbm.get_all_profiles(session)
+        self.assertEqual(2, session.invoke_api.call_count)
+
+    def _create_profile(self, profile_id, name):
+        profile = mock.Mock()
+        profile.profileId = profile_id
+        profile.name = name
+        return profile
+
+    @mock.patch.object(pbm, 'get_all_profiles')
+    def test_get_profile_id_by_name(self, get_all_profiles):
+        profiles = [self._create_profile(str(i), 'profile-%d' % i)
+                    for i in range(0, 10)]
+        get_all_profiles.return_value = profiles
+
+        session = mock.Mock()
+        exp_profile_id = '5'
+        profile_id = pbm.get_profile_id_by_name(session,
+                                                'profile-%s' % exp_profile_id)
+        self.assertEqual(exp_profile_id, profile_id)
+        get_all_profiles.assert_called_once_with(session)
+
+    @mock.patch.object(pbm, 'get_all_profiles')
+    def test_get_profile_id_by_name_with_invalid_profile(self,
+                                                         get_all_profiles):
+        profiles = [self._create_profile(str(i), 'profile-%d' % i)
+                    for i in range(0, 10)]
+        get_all_profiles.return_value = profiles
+
+        session = mock.Mock()
+        profile_id = pbm.get_profile_id_by_name(session,
+                                                'profile-%s' % (i + 1))
+        self.assertFalse(profile_id)
+        get_all_profiles.assert_called_once_with(session)
+
+    def test_filter_hubs_by_profile(self):
+        pbm_client = mock.Mock()
+        session = mock.Mock()
+        session.pbm = pbm_client
+        hubs = mock.Mock()
+        profile_id = 'profile-0'
+
+        pbm.filter_hubs_by_profile(session, hubs, profile_id)
+        session.invoke_api.assert_called_once_with(
+            pbm_client,
+            'PbmQueryMatchingHub',
+            pbm_client.service_content.placementSolver,
+            hubsToSearch=hubs,
+            profile=profile_id)
+
+    def _create_datastore(self, value):
+        ds = mock.Mock()
+        ds.value = value
+        return ds
+
+    def test_convert_datastores_to_hubs(self):
+        ds_values = []
+        datastores = []
+        for i in range(0, 10):
+            value = "ds-%d" % i
+            ds_values.append(value)
+            datastores.append(self._create_datastore(value))
+
+        pbm_client_factory = mock.Mock()
+        pbm_client_factory.create.side_effect = lambda *args: mock.Mock()
+        hubs = pbm.convert_datastores_to_hubs(pbm_client_factory, datastores)
+        self.assertEqual(len(datastores), len(hubs))
+        hub_ids = [hub.hubId for hub in hubs]
+        self.assertEqual(set(ds_values), set(hub_ids))
+
+    def test_filter_datastores_by_hubs(self):
+        ds_values = []
+        datastores = []
+        for i in range(0, 10):
+            value = "ds-%d" % i
+            ds_values.append(value)
+            datastores.append(self._create_datastore(value))
+
+        hubs = []
+        hub_ids = ds_values[0:len(ds_values) / 2]
+        for hub_id in hub_ids:
+            hub = mock.Mock()
+            hub.hubId = hub_id
+            hubs.append(hub)
+
+        filtered_ds = pbm.filter_datastores_by_hubs(hubs, datastores)
+        self.assertEqual(len(hubs), len(filtered_ds))
+        filtered_ds_values = [ds.value for ds in filtered_ds]
+        self.assertEqual(set(hub_ids), set(filtered_ds_values))