gh-80958: unittest: discovery support for namespace packages as start directory (#123820)

1 files changed, 52 insertions, 2 deletions

diff --git a/Lib/test/test_unittest/test_discovery.py b/Lib/test/test_unittest/test_discovery.py
index a44b18406c0..38c9779daaf 100644
--- a/Lib/test/test_unittest/test_discovery.py
+++ b/Lib/test/test_unittest/test_discovery.py
@@ -4,12 +4,14 @@ import re
 import sys
 import types
 import pickle
+from importlib._bootstrap_external import NamespaceLoader
 from test import support
 from test.support import import_helper
 
 import unittest
 import unittest.mock
 import test.test_unittest
+from test.test_importlib import util as test_util
 
 
 class TestableTestProgram(unittest.TestProgram):
@@ -395,7 +397,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
@@ -815,7 +817,7 @@ class TestDiscovery(unittest.TestCase):
         expectedPath = os.path.abspath(os.path.dirname(test.test_unittest.__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
@@ -848,6 +850,54 @@ class TestDiscovery(unittest.TestCase):
                          'Can not use builtin modules '
                          'as dotted module names')
 
+    def test_discovery_from_dotted_namespace_packages(self):
+        loader = unittest.TestLoader()
+
+        package = types.ModuleType('package')
+        package.__name__ = "tests"
+        package.__path__ = ['/a', '/b']
+        package.__file__ = None
+        package.__spec__ = types.SimpleNamespace(
+            name=package.__name__,
+            loader=NamespaceLoader(package.__name__, package.__path__, None),
+            submodule_search_locations=['/a', '/b']
+        )
+
+        def _import(packagename, *args, **kwargs):
+            sys.modules[packagename] = package
+            return package
+
+        _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
+
+        with unittest.mock.patch('builtins.__import__', _import):
+            # Since loader.discover() can modify sys.path, restore it when done.
+            with import_helper.DirsOnSysPath():
+                # Make sure to remove 'package' from sys.modules when done.
+                with test_util.uncache('package'):
+                    suite = loader.discover('package')
+
+        self.assertEqual(suite, ['/a/tests', '/b/tests'])
+
+    def test_discovery_start_dir_is_namespace(self):
+        """Subdirectory discovery not affected if start_dir is a namespace pkg."""
+        loader = unittest.TestLoader()
+        with (
+            import_helper.DirsOnSysPath(os.path.join(os.path.dirname(__file__))),
+            test_util.uncache('namespace_test_pkg')
+        ):
+            suite = loader.discover('namespace_test_pkg')
+        self.assertEqual(
+            {list(suite)[0]._tests[0].__module__ for suite in suite._tests if list(suite)},
+            # files under namespace_test_pkg.noop not discovered.
+            {'namespace_test_pkg.test_foo', 'namespace_test_pkg.bar.test_bar'},
+        )
+
     def test_discovery_failed_discovery(self):
         from test.test_importlib import util