1"""Tools not exempt from being descended into in tracebacks"""
2
3import time
4
5
6__all__ = ['make_decorator', 'raises', 'set_trace', 'timed', 'with_setup',
7           'TimeExpired', 'istest', 'nottest']
8
9
10class TimeExpired(AssertionError):
11    pass
12
13
14def make_decorator(func):
15    """
16    Wraps a test decorator so as to properly replicate metadata
17    of the decorated function, including nose's additional stuff
18    (namely, setup and teardown).
19    """
20    def decorate(newfunc):
21        if hasattr(func, 'compat_func_name'):
22            name = func.compat_func_name
23        else:
24            name = func.__name__
25        newfunc.__dict__ = func.__dict__
26        newfunc.__doc__ = func.__doc__
27        newfunc.__module__ = func.__module__
28        if not hasattr(newfunc, 'compat_co_firstlineno'):
29            newfunc.compat_co_firstlineno = func.func_code.co_firstlineno
30        try:
31            newfunc.__name__ = name
32        except TypeError:
33            # can't set func name in 2.3
34            newfunc.compat_func_name = name
35        return newfunc
36    return decorate
37
38
39def raises(*exceptions):
40    """Test must raise one of expected exceptions to pass.
41
42    Example use::
43
44      @raises(TypeError, ValueError)
45      def test_raises_type_error():
46          raise TypeError("This test passes")
47
48      @raises(Exception)
49      def test_that_fails_by_passing():
50          pass
51
52    If you want to test many assertions about exceptions in a single test,
53    you may want to use `assert_raises` instead.
54    """
55    valid = ' or '.join([e.__name__ for e in exceptions])
56    def decorate(func):
57        name = func.__name__
58        def newfunc(*arg, **kw):
59            try:
60                func(*arg, **kw)
61            except exceptions:
62                pass
63            except:
64                raise
65            else:
66                message = "%s() did not raise %s" % (name, valid)
67                raise AssertionError(message)
68        newfunc = make_decorator(func)(newfunc)
69        return newfunc
70    return decorate
71
72
73def set_trace():
74    """Call pdb.set_trace in the calling frame, first restoring
75    sys.stdout to the real output stream. Note that sys.stdout is NOT
76    reset to whatever it was before the call once pdb is done!
77    """
78    import pdb
79    import sys
80    stdout = sys.stdout
81    sys.stdout = sys.__stdout__
82    pdb.Pdb().set_trace(sys._getframe().f_back)
83
84
85def timed(limit):
86    """Test must finish within specified time limit to pass.
87
88    Example use::
89
90      @timed(.1)
91      def test_that_fails():
92          time.sleep(.2)
93    """
94    def decorate(func):
95        def newfunc(*arg, **kw):
96            start = time.time()
97            result = func(*arg, **kw)
98            end = time.time()
99            if end - start > limit:
100                raise TimeExpired("Time limit (%s) exceeded" % limit)
101            return result
102        newfunc = make_decorator(func)(newfunc)
103        return newfunc
104    return decorate
105
106
107def with_setup(setup=None, teardown=None):
108    """Decorator to add setup and/or teardown methods to a test function::
109
110      @with_setup(setup, teardown)
111      def test_something():
112          " ... "
113
114    Note that `with_setup` is useful *only* for test functions, not for test
115    methods or inside of TestCase subclasses.
116    """
117    def decorate(func, setup=setup, teardown=teardown):
118        if setup:
119            if hasattr(func, 'setup'):
120                _old_s = func.setup
121                def _s():
122                    setup()
123                    _old_s()
124                func.setup = _s
125            else:
126                func.setup = setup
127        if teardown:
128            if hasattr(func, 'teardown'):
129                _old_t = func.teardown
130                def _t():
131                    _old_t()
132                    teardown()
133                func.teardown = _t
134            else:
135                func.teardown = teardown
136        return func
137    return decorate
138
139
140def istest(func):
141    """Decorator to mark a function or method as a test
142    """
143    func.__test__ = True
144    return func
145
146
147def nottest(func):
148    """Decorator to mark a function or method as *not* a test
149    """
150    func.__test__ = False
151    return func
152