Issue 17457: extend test discovery to support namespace packages

2 files changed, 129 insertions, 11 deletions

diff --git a/Lib/unittest/loader.py b/Lib/unittest/loader.py
index e872fcc..808c50e 100644
--- a/Lib/unittest/loader.py
+++ b/Lib/unittest/loader.py
@@ -61,8 +61,9 @@ class TestLoader(object):
     def loadTestsFromTestCase(self, testCaseClass):
         """Return a suite of all tests cases contained in testCaseClass"""
         if issubclass(testCaseClass, suite.TestSuite):
-            raise TypeError("Test cases should not be derived from TestSuite." \
-                                " Maybe you meant to derive from TestCase?")
+            raise TypeError("Test cases should not be derived from "
+                            "TestSuite. Maybe you meant to derive from "
+                            "TestCase?")
         testCaseNames = self.getTestCaseNames(testCaseClass)
         if not testCaseNames and hasattr(testCaseClass, 'runTest'):
             testCaseNames = ['runTest']
@@ -200,6 +201,8 @@ class TestLoader(object):
         self._top_level_dir = top_level_dir
 
         is_not_importable = False
+        is_namespace = False
+        tests = []
         if os.path.isdir(os.path.abspath(start_dir)):
             start_dir = os.path.abspath(start_dir)
             if start_dir != top_level_dir:
@@ -213,15 +216,52 @@ class TestLoader(object):
             else:
                 the_module = sys.modules[start_dir]
                 top_part = start_dir.split('.')[0]
-                start_dir = os.path.abspath(os.path.dirname((the_module.__file__)))
+                try:
+                    start_dir = os.path.abspath(
+                       os.path.dirname((the_module.__file__)))
+                except AttributeError:
+                    # look for namespace packages
+                    try:
+                        spec = the_module.__spec__
+                    except AttributeError:
+                        spec = None
+
+                    if spec and spec.loader is None:
+                        if spec.submodule_search_locations is not None:
+                            is_namespace = True
+
+                            for path in the_module.__path__:
+                                if (not set_implicit_top and
+                                    not path.startswith(top_level_dir)):
+                                    continue
+                                self._top_level_dir = \
+                                    (path.split(the_module.__name__
+                                         .replace(".", os.path.sep))[0])
+                                tests.extend(self._find_tests(path,
+                                                              pattern,
+                                                              namespace=True))
+                    elif the_module.__name__ in sys.builtin_module_names:
+                        # builtin module
+                        raise TypeError('Can not use builtin modules '
+                                        'as dotted module names') from None
+                    else:
+                        raise TypeError(
+                            'don\'t know how to discover from {!r}'
+                            .format(the_module)) from None
+
                 if set_implicit_top:
-                    self._top_level_dir = self._get_directory_containing_module(top_part)
-                    sys.path.remove(top_level_dir)
+                    if not is_namespace:
+                        self._top_level_dir = \
+                           self._get_directory_containing_module(top_part)
+                        sys.path.remove(top_level_dir)
+                    else:
+                        sys.path.remove(top_level_dir)
 
         if is_not_importable:
             raise ImportError('Start directory is not importable: %r' % start_dir)
 
-        tests = list(self._find_tests(start_dir, pattern))
+        if not is_namespace:
+            tests = list(self._find_tests(start_dir, pattern))
         return self.suiteClass(tests)
 
     def _get_directory_containing_module(self, module_name):
@@ -254,7 +294,7 @@ class TestLoader(object):
         # override this method to use alternative matching strategy
         return fnmatch(path, pattern)
 
-    def _find_tests(self, start_dir, pattern):
+    def _find_tests(self, start_dir, pattern, namespace=False):
         """Used by discovery. Yields test suites it loads."""
         paths = sorted(os.listdir(start_dir))
 
@@ -287,7 +327,8 @@ class TestLoader(object):
                         raise ImportError(msg % (mod_name, module_dir, expected_dir))
                     yield self.loadTestsFromModule(module)
             elif os.path.isdir(full_path):
-                if not os.path.isfile(os.path.join(full_path, '__init__.py')):
+                if (not namespace and
+                    not os.path.isfile(os.path.join(full_path, '__init__.py'))):
                     continue
 
                 load_tests = None
@@ -304,7 +345,8 @@ class TestLoader(object):
                         # tests loaded from package file
                         yield tests
                     # recurse into the package
-                    yield from self._find_tests(full_path, pattern)
+                    yield from self._find_tests(full_path, pattern,
+                                                namespace=namespace)
                 else:
                     try:
                         yield load_tests(self, tests, pattern)
diff --git a/Lib/unittest/test/test_discovery.py b/Lib/unittest/test/test_discovery.py
index d4eff40..6b7b128 100644
--- a/Lib/unittest/test/test_discovery.py
+++ b/Lib/unittest/test/test_discovery.py
@@ -1,6 +1,8 @@
 import os
 import re
 import sys
+import types
+import builtins
 from test import support
 
 import unittest
@@ -173,7 +175,7 @@ class TestDiscovery(unittest.TestCase):
         self.addCleanup(restore_isdir)
 
         _find_tests_args = []
-        def _find_tests(start_dir, pattern):
+        def _find_tests(start_dir, pattern, namespace=None):
             _find_tests_args.append((start_dir, pattern))
             return ['tests']
         loader._find_tests = _find_tests
@@ -436,7 +438,7 @@ class TestDiscovery(unittest.TestCase):
         expectedPath = os.path.abspath(os.path.dirname(unittest.test.__file__))
 
         self.wasRun = False
-        def _find_tests(start_dir, pattern):
+        def _find_tests(start_dir, pattern, namespace=None):
             self.wasRun = True
             self.assertEqual(start_dir, expectedPath)
             return tests
@@ -446,5 +448,79 @@ class TestDiscovery(unittest.TestCase):
         self.assertEqual(suite._tests, tests)
 
 
+    def test_discovery_from_dotted_path_builtin_modules(self):
+
+        loader = unittest.TestLoader()
+
+        listdir = os.listdir
+        os.listdir = lambda _: ['test_this_does_not_exist.py']
+        isfile = os.path.isfile
+        isdir = os.path.isdir
+        os.path.isdir = lambda _: False
+        orig_sys_path = sys.path[:]
+        def restore():
+            os.path.isfile = isfile
+            os.path.isdir = isdir
+            os.listdir = listdir
+            sys.path[:] = orig_sys_path
+        self.addCleanup(restore)
+
+        with self.assertRaises(TypeError) as cm:
+            loader.discover('sys')
+        self.assertEqual(str(cm.exception),
+                         'Can not use builtin modules '
+                         'as dotted module names')
+
+    def test_discovery_from_dotted_namespace_packages(self):
+        loader = unittest.TestLoader()
+
+        orig_import = __import__
+        package = types.ModuleType('package')
+        package.__path__ = ['/a', '/b']
+        package.__spec__ = types.SimpleNamespace(
+           loader=None,
+           submodule_search_locations=['/a', '/b']
+        )
+
+        def _import(packagename, *args, **kwargs):
+            sys.modules[packagename] = package
+            return package
+
+        def cleanup():
+            builtins.__import__ = orig_import
+        self.addCleanup(cleanup)
+        builtins.__import__ = _import
+
+        _find_tests_args = []
+        def _find_tests(start_dir, pattern, namespace=None):
+            _find_tests_args.append((start_dir, pattern))
+            return ['%s/tests' % start_dir]
+
+        loader._find_tests = _find_tests
+        loader.suiteClass = list
+        suite = loader.discover('package')
+        self.assertEqual(suite, ['/a/tests', '/b/tests'])
+
+    def test_discovery_failed_discovery(self):
+        loader = unittest.TestLoader()
+        package = types.ModuleType('package')
+        orig_import = __import__
+
+        def _import(packagename, *args, **kwargs):
+            sys.modules[packagename] = package
+            return package
+
+        def cleanup():
+            builtins.__import__ = orig_import
+        self.addCleanup(cleanup)
+        builtins.__import__ = _import
+
+        with self.assertRaises(TypeError) as cm:
+            loader.discover('package')
+        self.assertEqual(str(cm.exception),
+                         'don\'t know how to discover from {!r}'
+                         .format(package))
+
+
 if __name__ == '__main__':
     unittest.main()