From 68868872195135bdb90d45a5cb0d609400943eae Mon Sep 17 00:00:00 2001
From: Andy Freeland <andy@andyfreeland.net>
Date: Tue, 10 Apr 2018 16:17:24 -0700
Subject: [PATCH] Add support for `assert_called()`

This was added in Python 3.6.

I also re-ordered the wrappers to match the order in the unittest.mock
docs to make it easier to see if any are missing in the future.

ref: https://docs.python.org/3/library/unittest.mock.html#unittest.mock.Mock.assert_called
---
 CHANGELOG.rst       |  7 +++++++
 pytest_mock.py      | 13 ++++++++++---
 test_pytest_mock.py | 12 ++++++++++++
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index d88da20..fbc198a 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,3 +1,10 @@
+NEXT
+----
+
+* Add support for the recently added ``assert_called`` method in Python 3.6 and ``mock-2.0``. Thanks `@rouge8`_ for the PR (`#115`_).
+
+.. _#115: https://github.com/pytest-dev/pytest-mock/pull/115
+
 1.9.0
 -----
 
diff --git a/pytest_mock.py b/pytest_mock.py
index db77e14..632598d 100644
--- a/pytest_mock.py
+++ b/pytest_mock.py
@@ -245,6 +245,12 @@ def wrap_assert_any_call(*args, **kwargs):
                    *args, **kwargs)
 
 
+def wrap_assert_called(*args, **kwargs):
+    __tracebackhide__ = True
+    assert_wrapper(_mock_module_originals["assert_called"],
+                   *args, **kwargs)
+
+
 def wrap_assert_methods(config):
     """
     Wrap assert methods of mock module so we can hide their traceback and
@@ -257,12 +263,13 @@ def wrap_assert_methods(config):
     mock_module = _get_mock_module(config)
 
     wrappers = {
-        'assert_not_called': wrap_assert_not_called,
+        'assert_called': wrap_assert_called,
+        'assert_called_once': wrap_assert_called_once,
         'assert_called_with': wrap_assert_called_with,
         'assert_called_once_with': wrap_assert_called_once_with,
-        'assert_called_once': wrap_assert_called_once,
-        'assert_has_calls': wrap_assert_has_calls,
         'assert_any_call': wrap_assert_any_call,
+        'assert_has_calls': wrap_assert_has_calls,
+        'assert_not_called': wrap_assert_not_called,
     }
     for method, wrapper in wrappers.items():
         try:
diff --git a/test_pytest_mock.py b/test_pytest_mock.py
index 906dbc9..d9a4faf 100644
--- a/test_pytest_mock.py
+++ b/test_pytest_mock.py
@@ -399,6 +399,18 @@ def test_assert_called_once_wrapper(mocker):
         stub.assert_called_once()
 
 
+def test_assert_called_wrapper(mocker):
+    stub = mocker.stub()
+    if not hasattr(stub, 'assert_called'):
+        pytest.skip('assert_called_once not available')
+    with assert_traceback():
+        stub.assert_called()
+    stub("foo")
+    stub.assert_called()
+    stub("foo")
+    stub.assert_called()
+
+
 @pytest.mark.usefixtures('needs_assert_rewrite')
 def test_assert_called_args_with_introspection(mocker):
     stub = mocker.stub()