Wrap assert methods of mock module to clean up traceback

pytest_mock.py

@@ -2,6 +2,7 @@
 import sys
 
 import pytest
+from distutils.util import strtobool
 
 if sys.version_info >= (3, 3): # pragma: no cover
     import unittest.mock as mock_module
@@ -137,3 +138,98 @@ def mock(mocker):
     warnings.warn('"mock" fixture has been deprecated, use "mocker" instead',
                   DeprecationWarning)
     return mocker
+
+
+_mock_module_patches = []
+_mock_module_originals = {}
+
+
+def assert_wrapper(method, *args, **kwargs):
+    __tracebackhide__ = True
+    try:
+        method(*args, **kwargs)
+    except AssertionError as e:
+        raise AssertionError(*e.args)
+
+
+def wrap_assert_not_called(*args, **kwargs):
+    __tracebackhide__ = True
+    assert_wrapper(_mock_module_originals["assert_not_called"],
+                   *args, **kwargs)
+
+
+def wrap_assert_called_with(*args, **kwargs):
+    __tracebackhide__ = True
+    assert_wrapper(_mock_module_originals["assert_called_with"],
+                   *args, **kwargs)
+
+
+def wrap_assert_called_once_with(*args, **kwargs):
+    __tracebackhide__ = True
+    assert_wrapper(_mock_module_originals["assert_called_once_with"],
+                   *args, **kwargs)
+
+
+def wrap_assert_has_calls(*args, **kwargs):
+    __tracebackhide__ = True
+    assert_wrapper(_mock_module_originals["assert_has_calls"],
+                   *args, **kwargs)
+
+
+def wrap_assert_any_call(*args, **kwargs):
+    __tracebackhide__ = True
+    assert_wrapper(_mock_module_originals["assert_any_call"],
+                   *args, **kwargs)
+
+
+def wrap_assert_methods(config):
+    """
+    Wrap assert methods of mock module so we can hide their traceback
+    """
+    # Make sure we only do this once
+    if _mock_module_originals:
+        return
+
+    wrappers = {
+        'assert_not_called': wrap_assert_not_called,
+        'assert_called_with': wrap_assert_called_with,
+        'assert_called_once_with': wrap_assert_called_once_with,
+        'assert_has_calls': wrap_assert_has_calls,
+        'assert_any_call': wrap_assert_any_call,
+    }
+    for method, wrapper in wrappers.items():
+        try:
+            original = getattr(mock_module.NonCallableMock, method)
+        except AttributeError:
+            continue
+        _mock_module_originals[method] = original
+        patcher = mock_module.patch.object(
+            mock_module.NonCallableMock, method, wrapper)
+        patcher.start()
+        _mock_module_patches.append(patcher)
+    config.add_cleanup(unwrap_assert_methods)
+
+
+def unwrap_assert_methods():
+    for patcher in _mock_module_patches:
+        patcher.stop()
+    _mock_module_patches[:] = []
+    _mock_module_originals.clear()
+
+
+def pytest_addoption(parser):
+    parser.addini('mock_traceback_monkeypatch',
+                  'Monkeypatch the mock library to improve reporting of the '
+                  'assert_called_... methods',
+                  default=True)
+
+
+def parse_ini_boolean(value):
+    if value in (True, False):
+        return value
+    return bool(strtobool(value))
+
+
+def pytest_configure(config):
+    if parse_ini_boolean(config.getini('mock_traceback_monkeypatch')):
+        wrap_assert_methods(config)

test_pytest_mock.py

@@ -1,6 +1,9 @@
 import os
 import platform
+import sys
+from contextlib import contextmanager
 
+import py.code
 import pytest
 
 
@@ -246,3 +249,91 @@ def bar(arg):
     assert Foo.bar(arg=10) == 20
     Foo.bar.assert_called_once_with(arg=10)
     spy.assert_called_once_with(arg=10)
+
+
+@contextmanager
+def assert_traceback():
+    """
+    Assert that this file is at the top of the filtered traceback
+    """
+    try:
+        yield
+    except AssertionError:
+        traceback = py.code.ExceptionInfo().traceback
+        crashentry = traceback.getcrashentry()
+        assert crashentry.path == __file__
+    else:
+        raise AssertionError("DID NOT RAISE")
+
+
+@pytest.mark.skipif(sys.version_info[:2] == (3, 4),
+                    reason="assert_not_called not available in python 3.4")
+def test_assert_not_called_wrapper(mocker):
+    stub = mocker.stub()
+    stub.assert_not_called()
+    stub()
+    with assert_traceback():
+        stub.assert_not_called()
+
+
+def test_assert_called_with_wrapper(mocker):
+    stub = mocker.stub()
+    stub("foo")
+    stub.assert_called_with("foo")
+    with assert_traceback():
+        stub.assert_called_with("bar")
+
+
+def test_assert_called_once_with_wrapper(mocker):
+    stub = mocker.stub()
+    stub("foo")
+    stub.assert_called_once_with("foo")
+    stub("foo")
+    with assert_traceback():
+        stub.assert_called_once_with("foo")
+
+
+def test_assert_any_call_wrapper(mocker):
+    stub = mocker.stub()
+    stub("foo")
+    stub("foo")
+    stub.assert_any_call("foo")
+    with assert_traceback():
+        stub.assert_any_call("bar")
+
+
+def test_assert_has_calls(mocker):
+    from pytest_mock import mock_module
+    stub = mocker.stub()
+    stub("foo")
+    stub.assert_has_calls([mock_module.call("foo")])
+    with assert_traceback():
+        stub.assert_has_calls([mock_module.call("bar")])
+
+
+def test_monkeypatch_ini(mocker, testdir):
+    # Make sure the following function actually tests something
+    stub = mocker.stub()
+    assert stub.assert_called_with.__module__ != stub.__module__
+
+    testdir.makepyfile("""
+        import py.code
+        def test_foo(mocker):
+            stub = mocker.stub()
+            assert stub.assert_called_with.__module__ == stub.__module__
+    """)
+    testdir.makeini("""
+        [pytest]
+        mock_traceback_monkeypatch = false
+    """)
+
+    result = testdir.runpytest_subprocess()
+    assert result.ret == 0
+
+
+def test_parse_ini_boolean(testdir):
+    import pytest_mock
+    assert pytest_mock.parse_ini_boolean('True') is True
+    assert pytest_mock.parse_ini_boolean('false') is False
+    with pytest.raises(ValueError):
+        pytest_mock.parse_ini_boolean('foo')