xref: /aosp_15_r20/external/grpc-grpc/src/python/grpcio_tests/tests/_loader.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2)
1# Copyright 2015 gRPC authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15from __future__ import absolute_import
16
17import importlib
18import logging
19import os
20import pkgutil
21import re
22import sys
23import unittest
24
25import coverage
26
27logger = logging.getLogger(__name__)
28
29TEST_MODULE_REGEX = r"^.*_test$"
30
31
32# Determines the path og a given path relative to the first matching
33# path on sys.path. Useful for determining what a directory's module
34# path will be.
35def _relativize_to_sys_path(path):
36    for sys_path in sys.path:
37        if path.startswith(sys_path):
38            relative = path[len(sys_path) :]
39            if not relative:
40                return ""
41            if relative.startswith(os.path.sep):
42                relative = relative[len(os.path.sep) :]
43            if not relative.endswith(os.path.sep):
44                relative += os.path.sep
45            return relative
46    raise AssertionError("Failed to relativize {} to sys.path.".format(path))
47
48
49def _relative_path_to_module_prefix(path):
50    return path.replace(os.path.sep, ".")
51
52
53class Loader(object):
54    """Test loader for setuptools test suite support.
55
56    Attributes:
57      suite (unittest.TestSuite): All tests collected by the loader.
58      loader (unittest.TestLoader): Standard Python unittest loader to be ran per
59        module discovered.
60      module_matcher (re.RegexObject): A regular expression object to match
61        against module names and determine whether or not the discovered module
62        contributes to the test suite.
63    """
64
65    def __init__(self):
66        self.suite = unittest.TestSuite()
67        self.loader = unittest.TestLoader()
68        self.module_matcher = re.compile(TEST_MODULE_REGEX)
69
70    def loadTestsFromNames(self, names, module=None):
71        """Function mirroring TestLoader::loadTestsFromNames, as expected by
72        setuptools.setup argument `test_loader`."""
73        # ensure that we capture decorators and definitions (else our coverage
74        # measure unnecessarily suffers)
75        coverage_context = coverage.Coverage(data_suffix=True)
76        coverage_context.start()
77        imported_modules = tuple(
78            importlib.import_module(name) for name in names
79        )
80        for imported_module in imported_modules:
81            self.visit_module(imported_module)
82        for imported_module in imported_modules:
83            try:
84                package_paths = imported_module.__path__
85            except AttributeError:
86                continue
87            self.walk_packages(package_paths)
88        coverage_context.stop()
89        coverage_context.save()
90        return self.suite
91
92    def walk_packages(self, package_paths):
93        """Walks over the packages, dispatching `visit_module` calls.
94
95        Args:
96          package_paths (list): A list of paths over which to walk through modules
97            along.
98        """
99        for path in package_paths:
100            self._walk_package(path)
101
102    def _walk_package(self, package_path):
103        prefix = _relative_path_to_module_prefix(
104            _relativize_to_sys_path(package_path)
105        )
106        for importer, module_name, is_package in pkgutil.walk_packages(
107            [package_path], prefix
108        ):
109            module = None
110            if module_name in sys.modules:
111                module = sys.modules[module_name]
112                self.visit_module(module)
113            else:
114                try:
115                    spec = importer.find_spec(module_name)
116                    module = importlib.util.module_from_spec(spec)
117                    spec.loader.exec_module(module)
118                    self.visit_module(module)
119                except ModuleNotFoundError:
120                    logger.debug("Skip loading %s", module_name)
121
122    def visit_module(self, module):
123        """Visits the module, adding discovered tests to the test suite.
124
125        Args:
126          module (module): Module to match against self.module_matcher; if matched
127            it has its tests loaded via self.loader into self.suite.
128        """
129        if self.module_matcher.match(module.__name__):
130            module_suite = self.loader.loadTestsFromModule(module)
131            self.suite.addTest(module_suite)
132
133
134def iterate_suite_cases(suite):
135    """Generator over all unittest.TestCases in a unittest.TestSuite.
136
137    Args:
138      suite (unittest.TestSuite): Suite to iterate over in the generator.
139
140    Returns:
141      generator: A generator over all unittest.TestCases in `suite`.
142    """
143    for item in suite:
144        if isinstance(item, unittest.TestSuite):
145            for child_item in iterate_suite_cases(item):
146                yield child_item
147        elif isinstance(item, unittest.TestCase):
148            yield item
149        else:
150            raise ValueError(
151                "unexpected suite item of type {}".format(type(item))
152            )
153