1import contextlib
2import re
3import sys
4
5
6def eq_(a, b, msg=None):
7    """Assert a == b, with repr messaging on failure."""
8    assert a == b, msg or "%r != %r" % (a, b)
9
10
11def ne_(a, b, msg=None):
12    """Assert a != b, with repr messaging on failure."""
13    assert a != b, msg or "%r == %r" % (a, b)
14
15
16def in_(a, b, msg=None):
17    """Assert a in b, with repr messaging on failure."""
18    assert a in b, msg or "%r not in %r" % (a, b)
19
20
21def not_in(a, b, msg=None):
22    """Assert a in not b, with repr messaging on failure."""
23    assert a not in b, msg or "%r is in %r" % (a, b)
24
25
26def _assert_proper_exception_context(exception):
27    """assert that any exception we're catching does not have a __context__
28    without a __cause__, and that __suppress_context__ is never set.
29
30    Python 3 will report nested as exceptions as "during the handling of
31    error X, error Y occurred". That's not what we want to do. We want
32    these exceptions in a cause chain.
33
34    """
35
36    if (
37        exception.__context__ is not exception.__cause__
38        and not exception.__suppress_context__
39    ):
40        assert False, (
41            "Exception %r was correctly raised but did not set a cause, "
42            "within context %r as its cause."
43            % (exception, exception.__context__)
44        )
45
46
47def _assert_proper_cause_cls(exception, cause_cls):
48    """assert that any exception we're catching does not have a __context__
49    without a __cause__, and that __suppress_context__ is never set.
50
51    Python 3 will report nested as exceptions as "during the handling of
52    error X, error Y occurred". That's not what we want to do. We want
53    these exceptions in a cause chain.
54
55    """
56    assert isinstance(exception.__cause__, cause_cls), (
57        "Exception %r was correctly raised but has cause %r, which does not "
58        "have the expected cause type %r."
59        % (exception, exception.__cause__, cause_cls)
60    )
61
62
63def assert_raises(except_cls, callable_, *args, **kw):
64    return _assert_raises(except_cls, callable_, args, kw)
65
66
67def assert_raises_with_proper_context(except_cls, callable_, *args, **kw):
68    return _assert_raises(except_cls, callable_, args, kw, check_context=True)
69
70
71def assert_raises_with_given_cause(
72    except_cls, cause_cls, callable_, *args, **kw
73):
74    return _assert_raises(except_cls, callable_, args, kw, cause_cls=cause_cls)
75
76
77def assert_raises_message(except_cls, msg, callable_, *args, **kwargs):
78    return _assert_raises(except_cls, callable_, args, kwargs, msg=msg)
79
80
81def assert_raises_message_with_proper_context(
82    except_cls, msg, callable_, *args, **kwargs
83):
84    return _assert_raises(
85        except_cls, callable_, args, kwargs, msg=msg, check_context=True
86    )
87
88
89def assert_raises_message_with_given_cause(
90    except_cls, msg, cause_cls, callable_, *args, **kwargs
91):
92    return _assert_raises(
93        except_cls, callable_, args, kwargs, msg=msg, cause_cls=cause_cls
94    )
95
96
97def _assert_raises(
98    except_cls,
99    callable_,
100    args,
101    kwargs,
102    msg=None,
103    check_context=False,
104    cause_cls=None,
105):
106    with _expect_raises(except_cls, msg, check_context, cause_cls) as ec:
107        callable_(*args, **kwargs)
108    return ec.error
109
110
111class _ErrorContainer:
112    error = None
113
114
115@contextlib.contextmanager
116def _expect_raises(except_cls, msg=None, check_context=False, cause_cls=None):
117    ec = _ErrorContainer()
118    if check_context:
119        are_we_already_in_a_traceback = sys.exc_info()[0]
120    try:
121        yield ec
122        success = False
123    except except_cls as err:
124        ec.error = err
125        success = True
126        if msg is not None:
127            # I'm often pdbing here, and "err" above isn't
128            # in scope, so assign the string explicitly
129            error_as_string = str(err)
130            assert re.search(msg, error_as_string, re.UNICODE), "%r !~ %s" % (
131                msg,
132                error_as_string,
133            )
134        if cause_cls is not None:
135            _assert_proper_cause_cls(err, cause_cls)
136        if check_context and not are_we_already_in_a_traceback:
137            _assert_proper_exception_context(err)
138        print(str(err).encode("utf-8"))
139
140    # it's generally a good idea to not carry traceback objects outside
141    # of the except: block, but in this case especially we seem to have
142    # hit some bug in either python 3.10.0b2 or greenlet or both which
143    # this seems to fix:
144    # https://github.com/python-greenlet/greenlet/issues/242
145    del ec
146
147    # assert outside the block so it works for AssertionError too !
148    assert success, "Callable did not raise an exception"
149
150
151def expect_raises(except_cls, check_context=False):
152    return _expect_raises(except_cls, check_context=check_context)
153
154
155def expect_raises_message(except_cls, msg, check_context=False):
156    return _expect_raises(except_cls, msg=msg, check_context=check_context)
157
158
159def expect_raises_with_proper_context(except_cls, check_context=True):
160    return _expect_raises(except_cls, check_context=check_context)
161
162
163def expect_raises_message_with_proper_context(
164    except_cls, msg, check_context=True
165):
166    return _expect_raises(except_cls, msg=msg, check_context=check_context)
167