mock: Add EventMock class

mariocj89 · tirkarthi · mariocj89 · commit 681abde2e457 · 2020-10-21T18:01:56.000+02:00
Add a new class that allows to wait for a call to happen by using
`Event` objects. This mock class can be used to test and validate
expectations of multithreading code.

Co-authored-by: Karthikeyan Singaravelan &lt;tir.karthi@gmail.com&gt;
diff --git a/Doc/library/unittest.mock.rst b/Doc/library/unittest.mock.rst
@@ -204,8 +204,10 @@ The Mock Class
     import asyncio
     import inspect
     import unittest
+    import threading
     from unittest.mock import sentinel, DEFAULT, ANY
     from unittest.mock import patch, call, Mock, MagicMock, PropertyMock, AsyncMock
+    from unittest.mock import ThreadingMock
     from unittest.mock import mock_open
 
 :class:`Mock` is a flexible mock object intended to replace the use of stubs and
@@ -1097,6 +1099,39 @@ object::
       [call('foo'), call('bar')]
 
 
+.. class:: ThreadingMock(spec=None, side_effect=None, return_value=DEFAULT, wraps=None, name=None, spec_set=None, unsafe=False, **kwargs)
+
+  A version of :class:`MagicMock` for multithreading tests. The
+  :class:`ThreadingMock` object provides extra methods to wait for a call to
+  happen on a different thread, rather than assert on it immediately.
+
+  .. method:: wait_until_called(mock_timeout=None)
+
+      Waits until the the mock is called.
+      If ``mock_timeout`` is set, after that number of seconds waiting,
+      it raises an :exc:`AssertionError`, waits forever otherwise.
+
+        >>> mock = ThreadingMock()
+        >>> thread = threading.Thread(target=mock)
+        >>> thread.start()
+        >>> mock.wait_until_called(mock_timeout=1)
+        >>> thread.join()
+
+  .. method:: wait_until_any_call(*args, mock_timeout=None, **kwargs)
+
+      Waits until the the mock is called with the specified arguments.
+      If ``mock_timeout`` is set, after that number of seconds waiting,
+      it raises an :exc:`AssertionError`, waits forever otherwise.
+
+        >>> mock = ThreadingMock()
+        >>> thread = threading.Thread(target=mock, args=(1,2,), kwargs={"arg": "thing"})
+        >>> thread.start()
+        >>> mock.wait_until_any_call(1, 2, arg="thing")
+        >>> thread.join()
+
+  .. versionadded:: 3.10
+
+
 Calling
 ~~~~~~~
 
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
@@ -14,6 +14,7 @@
     'call',
     'create_autospec',
     'AsyncMock',
+    'ThreadingMock',
     'FILTER_DIR',
     'NonCallableMock',
     'NonCallableMagicMock',
@@ -31,6 +32,7 @@
 import sys
 import builtins
 from asyncio import iscoroutinefunction
+import threading
 from types import CodeType, ModuleType, MethodType
 from unittest.util import safe_repr
 from functools import wraps, partial
@@ -2851,6 +2853,59 @@ def __set__(self, obj, val):
         self(val)
 
 
+class ThreadingMock(MagicMock):
+    """
+    A mock that can be used to wait until on calls happening
+    in a different thread.
+    """
+
+    def __init__(self, *args, **kwargs):
+        _safe_super(ThreadingMock, self).__init__(*args, **kwargs)
+        self.__dict__["_event"] = threading.Event()
+        self.__dict__["_expected_calls"] = []
+        self.__dict__["_events_lock"] = threading.Lock()
+
+    def __get_event(self, expected_args, expected_kwargs):
+        with self._events_lock:
+            for args, kwargs, event in self._expected_calls:
+                if (args, kwargs) == (expected_args, expected_kwargs):
+                    return event
+            new_event = threading.Event()
+            self._expected_calls.append((expected_args, expected_kwargs, new_event))
+            return new_event
+
+
+    def _mock_call(self, *args, **kwargs):
+        ret_value = _safe_super(ThreadingMock, self)._mock_call(*args, **kwargs)
+
+        call_event = self.__get_event(args, kwargs)
+        call_event.set()
+
+        self._event.set()
+
+        return ret_value
+
+    def wait_until_called(self, mock_timeout=None):
+        """Wait until the mock object is called.
+
+        `mock_timeout` - time to wait for in seconds, waits forever otherwise.
+        """
+        if not self._event.wait(timeout=mock_timeout):
+            msg = (f"{self._mock_name or 'mock'} was not called before"
+                   f" timeout({mock_timeout}).")
+            raise AssertionError(msg)
+
+    def wait_until_any_call(self, *args, mock_timeout=None, **kwargs):
+        """Wait until the mock object is called with given args.
+
+        `mock_timeout` - time to wait for in seconds, waits forever otherwise.
+        """
+        event = self.__get_event(args, kwargs)
+        if not event.wait(timeout=mock_timeout):
+            expected_string = self._format_mock_call_signature(args, kwargs)
+            raise AssertionError(f'{expected_string} call not found')
+
+
 def seal(mock):
     """Disable the automatic generation of child mocks.
 
diff --git a/Lib/unittest/test/testmock/testthreadingmock.py b/Lib/unittest/test/testmock/testthreadingmock.py
@@ -0,0 +1,170 @@
+import time
+import unittest
+import concurrent.futures
+
+from unittest.mock import patch, ThreadingMock, call
+
+
+class Something:
+
+    def method_1(self):
+        pass
+
+    def method_2(self):
+        pass
+
+
+class TestThreadingMock(unittest.TestCase):
+
+    def _call_after_delay(self, func, /, *args, **kwargs):
+        time.sleep(kwargs.pop('delay'))
+        func(*args, **kwargs)
+
+
+    def setUp(self):
+        self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=5)
+
+    def tearDown(self):
+        self._executor.shutdown()
+
+    def run_async(self, func, /, *args, delay=0, **kwargs):
+        self._executor.submit(self._call_after_delay, func, *args, **kwargs, delay=delay)
+
+    def _make_mock(self, *args, **kwargs):
+        return ThreadingMock(*args, **kwargs)
+
+    def test_instance_check(self):
+        waitable_mock = self._make_mock()
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+
+            self.assertIsInstance(something.method_1, ThreadingMock)
+            self.assertIsInstance(
+                something.method_1().method_2(), ThreadingMock)
+
+
+    def test_side_effect(self):
+        waitable_mock = self._make_mock()
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            something.method_1.side_effect = [1]
+
+            self.assertEqual(something.method_1(), 1)
+
+
+    def test_spec(self):
+        waitable_mock = self._make_mock(spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock) as m:
+            something = m()
+
+            self.assertIsInstance(something.method_1, ThreadingMock)
+            self.assertIsInstance(
+                something.method_1().method_2(), ThreadingMock)
+
+            with self.assertRaises(AttributeError):
+                m.test
+
+
+    def test_wait_until_called(self):
+        waitable_mock = self._make_mock(spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            self.run_async(something.method_1, delay=0.01)
+            something.method_1.wait_until_called()
+            something.method_1.assert_called_once()
+
+
+    def test_wait_until_called_called_before(self):
+        waitable_mock = self._make_mock(spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            something.method_1()
+            something.method_1.wait_until_called()
+            something.method_1.assert_called_once()
+
+
+    def test_wait_until_called_magic_method(self):
+        waitable_mock = self._make_mock(spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            self.run_async(something.method_1.__str__, delay=0.01)
+            something.method_1.__str__.wait_until_called()
+            something.method_1.__str__.assert_called_once()
+
+
+    def test_wait_until_called_timeout(self):
+        waitable_mock = self._make_mock(spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            self.run_async(something.method_1, delay=0.2)
+            with self.assertRaises(AssertionError):
+                something.method_1.wait_until_called(mock_timeout=0.1)
+            something.method_1.assert_not_called()
+
+            something.method_1.wait_until_called()
+            something.method_1.assert_called_once()
+
+
+    def test_wait_until_any_call_positional(self):
+        waitable_mock = self._make_mock(spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            self.run_async(something.method_1, 1, delay=0.1)
+            self.run_async(something.method_1, 2, delay=0.2)
+            self.run_async(something.method_1, 3, delay=0.3)
+            self.assertNotIn(call(1), something.method_1.mock_calls)
+
+            something.method_1.wait_until_any_call(1)
+            something.method_1.assert_called_once_with(1)
+            self.assertNotIn(call(2), something.method_1.mock_calls)
+            self.assertNotIn(call(3), something.method_1.mock_calls)
+
+            something.method_1.wait_until_any_call(3)
+            self.assertIn(call(2), something.method_1.mock_calls)
+            something.method_1.wait_until_any_call(2)
+
+
+    def test_wait_until_any_call_keywords(self):
+        waitable_mock = self._make_mock(spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            self.run_async(something.method_1, a=1, delay=0.1)
+            self.run_async(something.method_1, b=2, delay=0.2)
+            self.run_async(something.method_1, c=3, delay=0.3)
+            self.assertNotIn(call(a=1), something.method_1.mock_calls)
+
+            something.method_1.wait_until_any_call(a=1)
+            something.method_1.assert_called_once_with(a=1)
+            self.assertNotIn(call(b=2), something.method_1.mock_calls)
+            self.assertNotIn(call(c=3), something.method_1.mock_calls)
+
+            something.method_1.wait_until_any_call(c=3)
+            self.assertIn(call(b=2), something.method_1.mock_calls)
+            something.method_1.wait_until_any_call(b=2)
+
+    def test_wait_until_any_call_no_argument(self):
+        waitable_mock = self._make_mock(spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            something.method_1(1)
+
+            something.method_1.assert_called_once_with(1)
+            with self.assertRaises(AssertionError):
+                something.method_1.wait_until_any_call(mock_timeout=0.1)
+
+            something.method_1()
+            something.method_1.wait_until_any_call()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2019-09-13-13-28-10.bpo-17013.NWcgE3.rst b/Misc/NEWS.d/next/Library/2019-09-13-13-28-10.bpo-17013.NWcgE3.rst
@@ -0,0 +1,3 @@
+Add `ThreadingMock` to :mod:`unittest.mock` that can be used to create
+Mock objects that can wait until they are called. Patch by Karthikeyan
+Singaravelan and Mario Corchero.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+Add `ThreadingMock` to :mod:`unittest.mock` that can be used to create
	`2`	`+Mock objects that can wait until they are called. Patch by Karthikeyan`
	`3`	`+Singaravelan and Mario Corchero.`