Add Mock.assert_called()

Issue #26323: Add assert_called() and assert_called_once() methods to
unittest.mock.Mock.

3 files changed, 42 insertions, 0 deletions

diff --git a/NEWS b/NEWS
index 2e13b8a..2120418 100644
--- a/NEWS
+++ b/NEWS
@@ -1,6 +1,9 @@
 Library
 -------
 
+- Issue #26323: Add Mock.assert_called() and Mock.assert_called_once()
+  methods to unittest.mock. Patch written by Amit Saha.
+
 - Issue #22138: Fix mock.patch behavior when patching descriptors. Restore
   original values after patching. Patch contributed by Sean McCully.
 
diff --git a/mock/mock.py b/mock/mock.py
index 80e1994..c776aff 100644
--- a/mock/mock.py
+++ b/mock/mock.py
@@ -911,6 +911,24 @@ class NonCallableMock(Base):
                    (self._mock_name or 'mock', self.call_count))
             raise AssertionError(msg)
 
+    def assert_called(_mock_self):
+        """assert that the mock was called at least once
+        """
+        self = _mock_self
+        if self.call_count == 0:
+            msg = ("Expected '%s' to have been called." %
+                   self._mock_name or 'mock')
+            raise AssertionError(msg)
+
+    def assert_called_once(_mock_self):
+        """assert that the mock was called only once.
+        """
+        self = _mock_self
+        if not self.call_count == 1:
+            msg = ("Expected '%s' to have been called once. Called %s times." %
+                   (self._mock_name or 'mock', self.call_count))
+            raise AssertionError(msg)
+
     def assert_called_with(_mock_self, *args, **kwargs):
         """assert that the mock was called with the specified arguments.
 
diff --git a/mock/tests/testmock.py b/mock/tests/testmock.py
index 8ec9a60..7511c23 100644
--- a/mock/tests/testmock.py
+++ b/mock/tests/testmock.py
@@ -1288,6 +1288,27 @@ class MockTest(unittest.TestCase):
         with self.assertRaises(AssertionError):
             m.hello.assert_not_called()
 
+    def test_assert_called(self):
+        m = Mock()
+        with self.assertRaises(AssertionError):
+            m.hello.assert_called()
+        m.hello()
+        m.hello.assert_called()
+
+        m.hello()
+        m.hello.assert_called()
+
+    def test_assert_called_once(self):
+        m = Mock()
+        with self.assertRaises(AssertionError):
+            m.hello.assert_called_once()
+        m.hello()
+        m.hello.assert_called_once()
+
+        m.hello()
+        with self.assertRaises(AssertionError):
+            m.hello.assert_called_once()
+
     #Issue21256 printout of keyword args should be in deterministic order
     def test_sorted_call_signature(self):
         m = Mock()