Thanks Mario for the feedback.

> I see you inherit from Mock, should it inherit from MagicMock?

yes, it can inherit from MagicMock to mock magic methods and to wait on their call.

I thought some more about waiting for function call with arguments. One idea would be to have a dictionary with args to function as key mapping to an event object and to set and wait on that event object. To have wait_until_called that is always listening on a per mock object and to have wait_until_called_with listening on event object specific to that argument. Below is a sample implementation and an example to show it working with wait_until_called and wait_until_called_with. wait_until_called_with is something difficult to model since it needs per call event object and also need to store relevant key mapping to event object that is currently args and doesn't support keyword arguments. But wait_until_called is little simpler waiting on a per mock event object.

I will open up a PR with some tests.


class WaitableMock(MagicMock):

    def __init__(self, *args, **kwargs):
        _safe_super(WaitableMock, self).__init__(*args, **kwargs)
        self.event_class = kwargs.pop('event_class', threading.Event)
        self.event = self.event_class()
        self.expected_calls = {}

    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.get(call.args)
            if event and not event.is_set():
                event.set()

        # Always set per mock event object to ensure the function is called for wait_until_called.
        self.event.set()

        return ret_value

    def wait_until_called(self, timeout=1):
        return self.event.wait(timeout=timeout)

    def wait_until_called_with(self, *args, timeout=1):
        # If there are args create a per argument list event object and if not wait for per mock event object.
        if args:
            if args not in self.expected_calls:
                event = self.event_class()
                self.expected_calls[args] = event
            else:
                event = self.expected_calls[args]
        else:
            event = self.event

        return event.is_set() or event.wait(timeout=timeout)



# Sample program to wait on arguments, magic methods and validating wraps

import multiprocessing
import threading
import time
from unittest.mock import WaitableMock, patch, call

def call_after_sometime(func, *args, delay=1):
    time.sleep(delay)
    func(*args)

def wraps(*args):
    return 1

def foo(*args):
    pass

def bar(*args):
    pass

with patch('__main__.foo', WaitableMock(event_class=threading.Event, wraps=wraps)):
    with patch('__main__.bar', WaitableMock(event_class=threading.Event)):
        # Test normal call
        threading.Thread(target=call_after_sometime, args=(foo, 1), kwargs={'delay': 1}).start()
        threading.Thread(target=call_after_sometime, args=(bar, 1), kwargs={'delay': 5}).start()
        print("foo called ", foo.wait_until_called(timeout=2))
        print("bar called ", bar.wait_until_called(timeout=2))
        foo.assert_called_once()
        bar.assert_not_called()

        # Test wraps works
        assert foo() == 1

        # Test magic method
        threading.Thread(target=call_after_sometime, args=(foo.__str__, ), kwargs={'delay': 1}).start()
        print("foo.__str__ called ", foo.__str__.wait_until_called(timeout=2))
        print("bar.__str__ called ", bar.__str__.wait_until_called(timeout=2))

        foo.reset_mock()
        bar.reset_mock()

        # Test waiting for arguments
        threading.Thread(target=call_after_sometime, args=(bar, 1), kwargs={'delay': 1}).start()
        threading.Thread(target=call_after_sometime, args=(bar, 2), kwargs={'delay': 5}).start()
        print("bar called with 1 ", bar.wait_until_called_with(1, timeout=2))
        print("bar called with 2 ", bar.wait_until_called_with(2, timeout=2))
        time.sleep(5)
        print("bar called with 2 ", bar.wait_until_called_with(2))


$ ./python.exe ../backups/bpo17013_mock.py
foo called  True
bar called  False
foo.__str__ called  True
bar.__str__ called  False
bar called with 1  True
bar called with 2  False
bar called with 2  True