1ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh"""Loading unittests."""
2ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
3ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehimport os
4ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehimport re
5ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehimport sys
6ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehimport traceback
7ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehimport types
8ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
9ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehfrom functools import cmp_to_key as _CmpToKey
10ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehfrom fnmatch import fnmatch
11ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
12ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehfrom . import case, suite
13ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
14ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh__unittest = True
15ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
16ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh# what about .pyc or .pyo (etc)
17ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh# we would need to avoid loading the same tests multiple times
18ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh# from '.py', '.pyc' *and* '.pyo'
19ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehVALID_MODULE_NAME = re.compile(r'[_a-z]\w*\.py$', re.IGNORECASE)
20ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
21ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
22ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehdef _make_failed_import_test(name, suiteClass):
23ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    message = 'Failed to import test module: %s\n%s' % (name, traceback.format_exc())
24ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    return _make_failed_test('ModuleImportFailure', name, ImportError(message),
25ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                             suiteClass)
26ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
27ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehdef _make_failed_load_tests(name, exception, suiteClass):
28ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    return _make_failed_test('LoadTestsFailure', name, exception, suiteClass)
29ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
30ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehdef _make_failed_test(classname, methodname, exception, suiteClass):
31ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def testFailure(self):
32ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        raise exception
33ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    attrs = {methodname: testFailure}
34ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    TestClass = type(classname, (case.TestCase,), attrs)
35ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    return suiteClass((TestClass(methodname),))
36ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
37ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
38ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehclass TestLoader(object):
39ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    """
40ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    This class is responsible for loading tests according to various criteria
41ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    and returning them wrapped in a TestSuite
42ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    """
43ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    testMethodPrefix = 'test'
44ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    sortTestMethodsUsing = cmp
45ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    suiteClass = suite.TestSuite
46ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    _top_level_dir = None
47ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
48ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def loadTestsFromTestCase(self, testCaseClass):
49ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Return a suite of all tests cases contained in testCaseClass"""
50ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if issubclass(testCaseClass, suite.TestSuite):
51ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            raise TypeError("Test cases should not be derived from TestSuite." \
52ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                                " Maybe you meant to derive from TestCase?")
53ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        testCaseNames = self.getTestCaseNames(testCaseClass)
54ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if not testCaseNames and hasattr(testCaseClass, 'runTest'):
55ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            testCaseNames = ['runTest']
56ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        loaded_suite = self.suiteClass(map(testCaseClass, testCaseNames))
57ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return loaded_suite
58ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
59ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def loadTestsFromModule(self, module, use_load_tests=True):
60ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Return a suite of all tests cases contained in the given module"""
61ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        tests = []
62ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        for name in dir(module):
63ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            obj = getattr(module, name)
64ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            if isinstance(obj, type) and issubclass(obj, case.TestCase):
65ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                tests.append(self.loadTestsFromTestCase(obj))
66ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
67ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        load_tests = getattr(module, 'load_tests', None)
68ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        tests = self.suiteClass(tests)
69ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if use_load_tests and load_tests is not None:
70ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            try:
71ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                return load_tests(self, tests, None)
72ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            except Exception, e:
73ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                return _make_failed_load_tests(module.__name__, e,
74ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                                               self.suiteClass)
75ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return tests
76ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
77ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def loadTestsFromName(self, name, module=None):
78ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Return a suite of all tests cases given a string specifier.
79ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
80ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        The name may resolve either to a module, a test case class, a
81ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        test method within a test case class, or a callable object which
82ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        returns a TestCase or TestSuite instance.
83ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
84ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        The method optionally resolves the names relative to a given module.
85ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
86ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        parts = name.split('.')
87ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if module is None:
88ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            parts_copy = parts[:]
89ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            while parts_copy:
90ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                try:
91ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    module = __import__('.'.join(parts_copy))
92ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    break
93ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                except ImportError:
94ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    del parts_copy[-1]
95ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    if not parts_copy:
96ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        raise
97ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            parts = parts[1:]
98ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        obj = module
99ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        for part in parts:
100ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            parent, obj = obj, getattr(obj, part)
101ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
102ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if isinstance(obj, types.ModuleType):
103ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            return self.loadTestsFromModule(obj)
104ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        elif isinstance(obj, type) and issubclass(obj, case.TestCase):
105ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            return self.loadTestsFromTestCase(obj)
106ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        elif (isinstance(obj, types.UnboundMethodType) and
107ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh              isinstance(parent, type) and
108ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh              issubclass(parent, case.TestCase)):
109ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            return self.suiteClass([parent(obj.__name__)])
110ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        elif isinstance(obj, suite.TestSuite):
111ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            return obj
112ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        elif hasattr(obj, '__call__'):
113ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            test = obj()
114ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            if isinstance(test, suite.TestSuite):
115ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                return test
116ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            elif isinstance(test, case.TestCase):
117ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                return self.suiteClass([test])
118ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            else:
119ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                raise TypeError("calling %s returned %s, not a test" %
120ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                                (obj, test))
121ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        else:
122ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            raise TypeError("don't know how to make test from: %s" % obj)
123ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
124ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def loadTestsFromNames(self, names, module=None):
125ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Return a suite of all tests cases found using the given sequence
126ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        of string specifiers. See 'loadTestsFromName()'.
127ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
128ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        suites = [self.loadTestsFromName(name, module) for name in names]
129ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return self.suiteClass(suites)
130ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
131ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def getTestCaseNames(self, testCaseClass):
132ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Return a sorted sequence of method names found within testCaseClass
133ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
134ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        def isTestMethod(attrname, testCaseClass=testCaseClass,
135ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                         prefix=self.testMethodPrefix):
136ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            return attrname.startswith(prefix) and \
137ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                hasattr(getattr(testCaseClass, attrname), '__call__')
138ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        testFnNames = filter(isTestMethod, dir(testCaseClass))
139ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if self.sortTestMethodsUsing:
140ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            testFnNames.sort(key=_CmpToKey(self.sortTestMethodsUsing))
141ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return testFnNames
142ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
143ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def discover(self, start_dir, pattern='test*.py', top_level_dir=None):
144ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Find and return all test modules from the specified start
145ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        directory, recursing into subdirectories to find them. Only test files
146ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        that match the pattern will be loaded. (Using shell style pattern
147ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        matching.)
148ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
149ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        All test modules must be importable from the top level of the project.
150ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        If the start directory is not the top level directory then the top
151ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        level directory must be specified separately.
152ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
153ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        If a test package name (directory with '__init__.py') matches the
154ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        pattern then the package will be checked for a 'load_tests' function. If
155ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        this exists then it will be called with loader, tests, pattern.
156ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
157ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        If load_tests exists then discovery does  *not* recurse into the package,
158ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        load_tests is responsible for loading all tests in the package.
159ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
160ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        The pattern is deliberately not stored as a loader attribute so that
161ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        packages can continue discovery themselves. top_level_dir is stored so
162ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        load_tests does not need to pass this argument in to loader.discover().
163ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
164ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        set_implicit_top = False
165ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if top_level_dir is None and self._top_level_dir is not None:
166ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # make top_level_dir optional if called from load_tests in a package
167ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            top_level_dir = self._top_level_dir
168ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        elif top_level_dir is None:
169ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            set_implicit_top = True
170ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            top_level_dir = start_dir
171ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
172ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        top_level_dir = os.path.abspath(top_level_dir)
173ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
174ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if not top_level_dir in sys.path:
175ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # all test modules must be importable from the top level directory
176ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # should we *unconditionally* put the start directory in first
177ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # in sys.path to minimise likelihood of conflicts between installed
178ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # modules and development versions?
179ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            sys.path.insert(0, top_level_dir)
180ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        self._top_level_dir = top_level_dir
181ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
182ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        is_not_importable = False
183ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if os.path.isdir(os.path.abspath(start_dir)):
184ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            start_dir = os.path.abspath(start_dir)
185ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            if start_dir != top_level_dir:
186ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                is_not_importable = not os.path.isfile(os.path.join(start_dir, '__init__.py'))
187ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        else:
188ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # support for discovery from dotted module names
189ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            try:
190ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                __import__(start_dir)
191ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            except ImportError:
192ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                is_not_importable = True
193ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            else:
194ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                the_module = sys.modules[start_dir]
195ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                top_part = start_dir.split('.')[0]
196ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                start_dir = os.path.abspath(os.path.dirname((the_module.__file__)))
197ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                if set_implicit_top:
198ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    self._top_level_dir = self._get_directory_containing_module(top_part)
199ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    sys.path.remove(top_level_dir)
200ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
201ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if is_not_importable:
202ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            raise ImportError('Start directory is not importable: %r' % start_dir)
203ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
204ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        tests = list(self._find_tests(start_dir, pattern))
205ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return self.suiteClass(tests)
206ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
207ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def _get_directory_containing_module(self, module_name):
208ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        module = sys.modules[module_name]
209ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        full_path = os.path.abspath(module.__file__)
210ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
211ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if os.path.basename(full_path).lower().startswith('__init__.py'):
212ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            return os.path.dirname(os.path.dirname(full_path))
213ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        else:
214ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # here we have been given a module rather than a package - so
215ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # all we can do is search the *same* directory the module is in
216ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            # should an exception be raised instead
217ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            return os.path.dirname(full_path)
218ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
219ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def _get_name_from_path(self, path):
220ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        path = os.path.splitext(os.path.normpath(path))[0]
221ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
222ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        _relpath = os.path.relpath(path, self._top_level_dir)
223ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        assert not os.path.isabs(_relpath), "Path must be within the project"
224ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        assert not _relpath.startswith('..'), "Path must be within the project"
225ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
226ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        name = _relpath.replace(os.path.sep, '.')
227ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return name
228ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
229ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def _get_module_from_name(self, name):
230ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        __import__(name)
231ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return sys.modules[name]
232ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
233ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def _match_path(self, path, full_path, pattern):
234ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        # override this method to use alternative matching strategy
235ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return fnmatch(path, pattern)
236ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
237ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def _find_tests(self, start_dir, pattern):
238ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Used by discovery. Yields test suites it loads."""
239ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        paths = os.listdir(start_dir)
240ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
241ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        for path in paths:
242ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            full_path = os.path.join(start_dir, path)
243ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            if os.path.isfile(full_path):
244ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                if not VALID_MODULE_NAME.match(path):
245ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    # valid Python identifiers only
246ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    continue
247ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                if not self._match_path(path, full_path, pattern):
248ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    continue
249ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                # if the test file matches, load it
250ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                name = self._get_name_from_path(full_path)
251ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                try:
252ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    module = self._get_module_from_name(name)
253ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                except:
254ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    yield _make_failed_import_test(name, self.suiteClass)
255ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                else:
256ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    mod_file = os.path.abspath(getattr(module, '__file__', full_path))
257ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    realpath = os.path.splitext(mod_file)[0]
258ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    fullpath_noext = os.path.splitext(full_path)[0]
259ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    if realpath.lower() != fullpath_noext.lower():
260ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        module_dir = os.path.dirname(realpath)
261ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        mod_name = os.path.splitext(os.path.basename(full_path))[0]
262ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        expected_dir = os.path.dirname(full_path)
263ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        msg = ("%r module incorrectly imported from %r. Expected %r. "
264ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                               "Is this module globally installed?")
265ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        raise ImportError(msg % (mod_name, module_dir, expected_dir))
266ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    yield self.loadTestsFromModule(module)
267ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            elif os.path.isdir(full_path):
268ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                if not os.path.isfile(os.path.join(full_path, '__init__.py')):
269ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    continue
270ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
271ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                load_tests = None
272ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                tests = None
273ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                if fnmatch(path, pattern):
274ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    # only check load_tests if the package directory itself matches the filter
275ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    name = self._get_name_from_path(full_path)
276ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    package = self._get_module_from_name(name)
277ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    load_tests = getattr(package, 'load_tests', None)
278ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    tests = self.loadTestsFromModule(package, use_load_tests=False)
279ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
280ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                if load_tests is None:
281ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    if tests is not None:
282ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        # tests loaded from package file
283ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        yield tests
284ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    # recurse into the package
285ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    for test in self._find_tests(full_path, pattern):
286ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        yield test
287ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                else:
288ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    try:
289ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        yield load_tests(self, tests, pattern)
290ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    except Exception, e:
291ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                        yield _make_failed_load_tests(package.__name__, e,
292ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                                                      self.suiteClass)
293ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
294ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehdefaultTestLoader = TestLoader()
295ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
296ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
297ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehdef _makeLoader(prefix, sortUsing, suiteClass=None):
298ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    loader = TestLoader()
299ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    loader.sortTestMethodsUsing = sortUsing
300ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    loader.testMethodPrefix = prefix
301ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    if suiteClass:
302ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        loader.suiteClass = suiteClass
303ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    return loader
304ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
305ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehdef getTestCaseNames(testCaseClass, prefix, sortUsing=cmp):
306ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    return _makeLoader(prefix, sortUsing).getTestCaseNames(testCaseClass)
307ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
308ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehdef makeSuite(testCaseClass, prefix='test', sortUsing=cmp,
309ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh              suiteClass=suite.TestSuite):
310ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    return _makeLoader(prefix, sortUsing, suiteClass).loadTestsFromTestCase(testCaseClass)
311ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
312ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehdef findTestCases(module, prefix='test', sortUsing=cmp,
313ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                  suiteClass=suite.TestSuite):
314ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    return _makeLoader(prefix, sortUsing, suiteClass).loadTestsFromModule(module)
315