diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index 8684f1dfa5729f..9aac9a57da8b65 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -18,6 +18,7 @@ 'NonCallableMagicMock', 'mock_open', 'PropertyMock', + 'WaitableMock', 'seal', ) @@ -25,9 +26,11 @@ __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,46 @@ 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 = 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. + """ + event = self._expected_calls[args] + return event.wait(timeout=timeout) + + def seal(mock): """Disable the automatic generation of child mocks. diff --git a/Lib/unittest/test/testmock/testwaitablemock.py b/Lib/unittest/test/testmock/testwaitablemock.py new file mode 100644 index 00000000000000..e929c70fdee4ed --- /dev/null +++ b/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() diff --git a/Misc/NEWS.d/next/Library/2019-04-13-20-56-09.bpo-17013.R_sgfy.rst b/Misc/NEWS.d/next/Library/2019-04-13-20-56-09.bpo-17013.R_sgfy.rst new file mode 100644 index 00000000000000..7b5afc1dde588d --- /dev/null +++ b/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.