[WIP] bpo-17013: Implement WaitableMock to create Mock objects that can wait until called

Lib/unittest/mock.py

@@ -18,16 +18,19 @@
     'NonCallableMagicMock',
     'mock_open',
     'PropertyMock',
+    'WaitableMock',
     'seal',
 )
 
 
 __version__ = '1.0'
 
 
+from collections import defaultdict
 import inspect
 import pprint
 import sys
+import threading
 import builtins
 from types import ModuleType
 from unittest.util import safe_repr
@@ -2482,6 +2485,52 @@ def __set__(self, obj, val):
         self(val)
 
 
+class WaitableMock(MagicMock):
+    """
+    A mock that can be used to wait until it was called.
+
+    `event_class` - Class to be used to create event object.
+    Defaults to Threading.Event and can take values like multiprocessing.Event.
+    """
+
+    def __init__(self, *args, event_class=threading.Event, **kwargs):
+        _safe_super(WaitableMock, self).__init__(*args, **kwargs)
+        self._event_class = event_class
+        self._event = event_class()
+        self._expected_calls = defaultdict(event_class)
+
+    def _mock_call(self, *args, **kwargs):
+        ret_value = _safe_super(WaitableMock, self)._mock_call(*args, **kwargs)
+
+        for call in self._mock_mock_calls:
+            event = self._expected_calls[call.args]
+            event.set()
+
+        self._event.set()
+
+        return ret_value
+
+    def wait_until_called(self, timeout=1.0):
+        """Wait until the mock object is called.
+
+        `timeout` - time to wait for in seconds. Defaults to 1.
+        """
+        return self._event.wait(timeout=timeout)
+
+    def wait_until_called_with(self, *args, timeout=1.0):
+        """Wait until the mock object is called with given args.
+
+        `timeout` - time to wait for in seconds. Defaults to 1.
+        """
+        if args not in self._expected_calls:
+            event = self._event_class()
+            self._expected_calls[args] = event
+        else:
+            event = self._expected_calls[args]
+
+        return event.wait(timeout=timeout)
+
+
 def seal(mock):
     """Disable the automatic generation of child mocks.
 

Lib/unittest/test/testmock/testwaitablemock.py

@@ -0,0 +1,144 @@
+import threading
+import time
+import unittest
+
+from test.support import start_threads
+from unittest.mock import patch, WaitableMock, call
+
+
+class Something:
+
+    def method_1(self):
+        pass
+
+    def method_2(self):
+        pass
+
+
+class TestWaitableMock(unittest.TestCase):
+
+
+    def _call_after_delay(self, func, *args, delay):
+        time.sleep(delay)
+        func(*args)
+
+
+    def _create_thread(self, func, *args, **kwargs):
+        thread = threading.Thread(target=self._call_after_delay,
+                                  args=(func,) + args, kwargs=kwargs)
+        return thread
+
+
+    def test_instance_check(self):
+        waitable_mock = WaitableMock(event_class=threading.Event)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+
+            self.assertIsInstance(something.method_1, WaitableMock)
+            self.assertIsInstance(
+                something.method_1().method_2(), WaitableMock)
+
+
+    def test_side_effect(self):
+        waitable_mock = WaitableMock(event_class=threading.Event)
+
+        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 = WaitableMock(
+            event_class=threading.Event, spec=Something)
+
+        with patch(f'{__name__}.Something', waitable_mock) as m:
+            something = m()
+
+            self.assertIsInstance(something.method_1, WaitableMock)
+            self.assertIsInstance(
+                something.method_1().method_2(), WaitableMock)
+
+            with self.assertRaises(AttributeError):
+                m.test
+
+
+    def test_wait_until_called(self):
+        waitable_mock = WaitableMock(event_class=threading.Event)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            thread = self._create_thread(something.method_1, delay=0.5)
+
+            with start_threads([thread]):
+                something.method_1.wait_until_called()
+                something.method_1.assert_called_once()
+
+
+    def test_wait_until_called_magic_method(self):
+        waitable_mock = WaitableMock(event_class=threading.Event)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            thread = self._create_thread(something.method_1.__str__, delay=0.5)
+
+            with start_threads([thread]):
+                something.method_1.__str__.wait_until_called()
+                something.method_1.__str__.assert_called_once()
+
+
+    def test_wait_until_called_timeout(self):
+        waitable_mock = WaitableMock(event_class=threading.Event)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            thread = self._create_thread(something.method_1, delay=0.5)
+
+            with start_threads([thread]):
+                something.method_1.wait_until_called(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_called_with(self):
+        waitable_mock = WaitableMock(event_class=threading.Event)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            thread_1 = self._create_thread(something.method_1, 1, delay=0.5)
+            thread_2 = self._create_thread(something.method_2, 1, delay=0.1)
+            thread_3 = self._create_thread(something.method_2, 2, delay=0.1)
+
+            with start_threads([thread_1, thread_2, thread_3]):
+                something.method_1.wait_until_called_with(1, timeout=0.1)
+                something.method_1.assert_not_called()
+
+                something.method_1.wait_until_called(timeout=2.0)
+                something.method_1.assert_called_once_with(1)
+                self.assertEqual(something.method_1.mock_calls, [call(1)])
+                something.method_2.assert_has_calls(
+                    [call(1), call(2)], any_order=True)
+
+
+    def test_wait_until_called_with_no_argument(self):
+        waitable_mock = WaitableMock(event_class=threading.Event)
+
+        with patch(f'{__name__}.Something', waitable_mock):
+            something = Something()
+            something.method_1(1)
+
+            something.method_1.assert_called_once_with(1)
+            self.assertFalse(
+                something.method_1.wait_until_called_with(timeout=0.1))
+
+            thread_1 = self._create_thread(something.method_1, delay=0.1)
+            with start_threads([thread_1]):
+                self.assertTrue(something.method_1.wait_until_called_with())
+
+
+if __name__ == "__main__":
+    unittest.main()

Misc/NEWS.d/next/Library/2019-04-13-20-56-09.bpo-17013.R_sgfy.rst

@@ -0,0 +1,3 @@
+Add `WaitableMock` to :mod:`unittest.mock` that can be used to create Mock
+objects that can wait until they are called. Patch by Karthikeyan
+Singaravelan.