Merge "import_modules_from_package: refactor with generators"

2 files changed, 20 insertions, 41 deletions

diff --git a/tests/unit/common/test_utils.py b/tests/unit/common/test_utils.py
index 663226242..923ec4aaa 100644
--- a/tests/unit/common/test_utils.py
+++ b/tests/unit/common/test_utils.py
@@ -20,6 +20,7 @@ from itertools import product, chain
 import mock
 from six.moves import configparser
 
+import yardstick
 from yardstick.common import utils
 from yardstick.common import constants
 
@@ -45,47 +46,25 @@ class IterSubclassesTestCase(unittest.TestCase):
         self.assertEqual([B, C, D], list(utils.itersubclasses(A)))
 
 
-class TryAppendModuleTestCase(unittest.TestCase):
-
-    @mock.patch('yardstick.common.utils.importutils')
-    def test_try_append_module_not_in_modules(self, mock_importutils):
-
-        modules = {}
-        name = 'foo'
-        utils.try_append_module(name, modules)
-        mock_importutils.import_module.assert_called_with(name)
-
-    @mock.patch('yardstick.common.utils.importutils')
-    def test_try_append_module_already_in_modules(self, mock_importutils):
-
-        modules = {'foo'}
-        name = 'foo'
-        utils.try_append_module(name, modules)
-        self.assertFalse(mock_importutils.import_module.called)
-
-
 class ImportModulesFromPackageTestCase(unittest.TestCase):
 
     @mock.patch('yardstick.common.utils.os.walk')
-    @mock.patch('yardstick.common.utils.try_append_module')
-    def test_import_modules_from_package_no_mod(self, mock_append, mock_walk):
-
-        sep = os.sep
+    def test_import_modules_from_package_no_mod(self, mock_walk):
+        yardstick_root = os.path.dirname(os.path.dirname(yardstick.__file__))
         mock_walk.return_value = ([
-            ('..' + sep + 'foo', ['bar'], ['__init__.py']),
-            ('..' + sep + 'foo' + sep + 'bar', [], ['baz.txt', 'qux.rst'])
+            (os.path.join(yardstick_root, 'foo'), ['bar'], ['__init__.py']),
+            (os.path.join(yardstick_root, 'foo', 'bar'), [], ['baz.txt', 'qux.rst'])
         ])
 
         utils.import_modules_from_package('foo.bar')
-        self.assertFalse(mock_append.called)
 
     @mock.patch('yardstick.common.utils.os.walk')
     @mock.patch('yardstick.common.utils.importutils')
     def test_import_modules_from_package(self, mock_importutils, mock_walk):
 
-        sep = os.sep
+        yardstick_root = os.path.dirname(os.path.dirname(yardstick.__file__))
         mock_walk.return_value = ([
-            ('foo' + sep + '..' + sep + 'bar', [], ['baz.py'])
+            (os.path.join(yardstick_root, 'foo', os.pardir, 'bar'), [], ['baz.py'])
         ])
 
         utils.import_modules_from_package('foo.bar')
diff --git a/yardstick/common/utils.py b/yardstick/common/utils.py
index c7ae9c1ef..1d7ea071c 100644
--- a/yardstick/common/utils.py
+++ b/yardstick/common/utils.py
@@ -70,26 +70,26 @@ def itersubclasses(cls, _seen=None):
                 yield sub
 
 
-def try_append_module(name, modules):
-    if name not in modules:
-        modules[name] = importutils.import_module(name)
-
-
 def import_modules_from_package(package):
     """Import modules from package and append into sys.modules
 
     :param: package - Full package name. For example: rally.deploy.engines
     """
-    path = [os.path.dirname(yardstick.__file__), ".."] + package.split(".")
-    path = os.path.join(*path)
+    yardstick_root = os.path.dirname(os.path.dirname(yardstick.__file__))
+    path = os.path.join(yardstick_root, *package.split("."))
     for root, dirs, files in os.walk(path):
-        for filename in files:
-            if filename.startswith("__") or not filename.endswith(".py"):
-                continue
-            new_package = ".".join(root.split(os.sep)).split("....")[1]
-            module_name = "%s.%s" % (new_package, filename[:-3])
+        matches = (filename for filename in files if filename.endswith(".py") and
+                   not filename.startswith("__"))
+        new_package = os.path.relpath(root, yardstick_root).replace(os.sep, ".")
+        module_names = set(
+            ("{}.{}".format(new_package, filename.rsplit(".py", 1)[0]) for filename in matches))
+        # find modules which haven't already been imported
+        missing_modules = module_names.difference(sys.modules)
+        logger.debug("importing %s", missing_modules)
+        # we have already checked for already imported modules, so we don't need to check again
+        for module_name in missing_modules:
             try:
-                try_append_module(module_name, sys.modules)
+                sys.modules[module_name] = importutils.import_module(module_name)
             except ImportError:
                 logger.exception("unable to import %s", module_name)