Add a way to discover only contrib extensions

Several OS projects(cinder, neutron, osc...) use
`novaclient.discover_extensions` for initialization novaclient.client.Client
with novaclient.v2.contrib extensions. In this case, it would be nice to
provide a way to not discover extension via python path an entry-point.

Change-Id: I030f4c55c2795c7f7973f5f12e54b9819c4a5578
Closes-Bug: #1509500

2 files changed, 67 insertions, 11 deletions

diff --git a/novaclient/client.py b/novaclient/client.py
index 00381524..579a315d 100644
--- a/novaclient/client.py
+++ b/novaclient/client.py
@@ -723,19 +723,24 @@ def _construct_http_client(username=None, password=None, project_id=None,
                           api_version=api_version)
 
 
-def discover_extensions(version):
+def discover_extensions(version, only_contrib=False):
+    """Returns the list of extensions, which can be discovered by python path,
+    contrib path and by entry-point 'novaclient.extension'.
+
+    :param version: api version
+    :type version: str or novaclient.api_versions.APIVersion
+    :param only_contrib: search only in contrib directory or not
+    :type only_contrib: bool
+    """
     if not isinstance(version, api_versions.APIVersion):
         version = api_versions.get_api_version(version)
-    extensions = []
-    for name, module in itertools.chain(
-            _discover_via_python_path(),
-            _discover_via_contrib_path(version),
-            _discover_via_entry_points()):
-
-        extension = ext.Extension(name, module)
-        extensions.append(extension)
-
-    return extensions
+    if only_contrib:
+        chain = _discover_via_contrib_path(version)
+    else:
+        chain = itertools.chain(_discover_via_python_path(),
+                                _discover_via_contrib_path(version),
+                                _discover_via_entry_points())
+    return [ext.Extension(name, module) for name, module in chain]
 
 
 def _discover_via_python_path():
diff --git a/novaclient/tests/unit/test_client.py b/novaclient/tests/unit/test_client.py
index f63fdd4a..1325bf72 100644
--- a/novaclient/tests/unit/test_client.py
+++ b/novaclient/tests/unit/test_client.py
@@ -22,6 +22,7 @@ from keystoneauth1 import adapter
 import mock
 import requests
 
+import novaclient.api_versions
 import novaclient.client
 import novaclient.extension
 from novaclient.tests.unit import utils
@@ -454,3 +455,53 @@ class SessionClientTest(utils.TestCase):
         client.request("http://no.where", 'GET')
         self.assertEqual(1, len(client.times))
         self.assertEqual('GET http://no.where', client.times[0][0])
+
+
+class DiscoverExtensionTest(utils.TestCase):
+
+    @mock.patch("novaclient.client._discover_via_entry_points")
+    @mock.patch("novaclient.client._discover_via_contrib_path")
+    @mock.patch("novaclient.client._discover_via_python_path")
+    @mock.patch("novaclient.extension.Extension")
+    def test_discover_all(self, mock_extension,
+                          mock_discover_via_python_path,
+                          mock_discover_via_contrib_path,
+                          mock_discover_via_entry_points):
+        def make_gen(start, end):
+            def f(*args, **kwargs):
+                for i in range(start, end):
+                    yield "name-%s" % i, i
+            return f
+
+        mock_discover_via_python_path.side_effect = make_gen(0, 3)
+        mock_discover_via_contrib_path.side_effect = make_gen(3, 5)
+        mock_discover_via_entry_points.side_effect = make_gen(5, 6)
+
+        version = novaclient.api_versions.APIVersion("2.0")
+
+        result = novaclient.client.discover_extensions(version)
+
+        self.assertEqual([mock.call("name-%s" % i, i) for i in range(0, 6)],
+                         mock_extension.call_args_list)
+        mock_discover_via_python_path.assert_called_once_with()
+        mock_discover_via_contrib_path.assert_called_once_with(version)
+        mock_discover_via_entry_points.assert_called_once_with()
+        self.assertEqual([mock_extension()] * 6, result)
+
+    @mock.patch("novaclient.client._discover_via_entry_points")
+    @mock.patch("novaclient.client._discover_via_contrib_path")
+    @mock.patch("novaclient.client._discover_via_python_path")
+    @mock.patch("novaclient.extension.Extension")
+    def test_discover_only_contrib(self, mock_extension,
+                                   mock_discover_via_python_path,
+                                   mock_discover_via_contrib_path,
+                                   mock_discover_via_entry_points):
+        mock_discover_via_contrib_path.return_value = [("name", "module")]
+
+        version = novaclient.api_versions.APIVersion("2.0")
+
+        novaclient.client.discover_extensions(version, only_contrib=True)
+        mock_discover_via_contrib_path.assert_called_once_with(version)
+        self.assertFalse(mock_discover_via_python_path.called)
+        self.assertFalse(mock_discover_via_entry_points.called)
+        mock_extension.assert_called_once_with("name", "module")