A flyweight mocking helper for Python

By
Posted on
Tags: python, testing, mocks

Recently, I needed a lightweight mocking solution when adding unit tests to some existing Python code. I could have used one of the many Python mocking libraries, but I only had to replace a few module functions during testing, so I just wrote a tiny mocking helper rather than add a dependency to the project.

The helper is simple and surprisingly versatile:

import contextlib
import functools

@contextlib.contextmanager
def mocked(func,
           expecting=None, returning=None, raising=None,  # specs
           replacement_logic=None, called=1):
    """Stub out and mock a function for a yield's duration."""

    if (returning, raising, replacement_logic).count(None) < 2:
        raise ValueError("returning, raising, and replacement_logic "
                         "are incompatible with each other")

    # default logic for implementing mock fn: return or raise per specs
    def default_logic(*_args, **_kwds):
        if raising:
            raise raising
        return returning

    # prepare wrapper to replace mocked function for duration of yield
    invocations = [0]
    @functools.wraps(func)
    def replacement(*args, **kwds):
        if expecting is not None:
            assert expecting == (args, kwds)  # did we get expected args?
        invocations[0] += 1
        return (replacement_logic or default_logic)(*args, **kwds)

    # replace mocked function, yield to test, and then check & clean up
    module = sys.modules.get(func.__module__)
    setattr(module, func.__name__, replacement)
    try:
        yield  # give control back to test for a while
        assert invocations[0] == called  # was mock called enough?
    finally:
        setattr(module, func.__name__, func)

def uncalled(func):
    """Require that a function not be called for a yield's duration."""
    return mocked(func, called=0)

The idea is to wrap test code that requires mocked external functions with a mocking helper. Here’s an example:

def test_ml_errors_must_be_reported(self):
    """When an error occurs, it must be reported; nothing must be sent."""
    data = self._request_data()
    exception = mls.MlError()
    exception.exc = '[[[error description]]]'
    with mocked(mls.subscriber_count, raising=exception):
        with uncalled(sms.send):
            resp = self.app.request("/mailings", method='POST', data=data)
    self.assertEqual(resp.status, '200 OK')
    self.assertIn(exception.exc, resp.data)  # must be reported

In this test method, I’m making a simulated POST request to a web service that’s supposed to check how many users are in a mailing list and then, possibly, send some SMS messages. In this case, however, I want to simulate that an error occurs when talking to the mailing-list service. The test must verify that the error is reported and, crucially, that no messages are sent.

The library modules for talking to the mailing-list service and the SMS service are called mls and sms. So, for the duration of the simulated request, I’m replacing two functions in these modules with mock versions. The mock version of mls.subscriber_count, when called, will raise the exception I’m trying to simulate. The mock version of sms.send, however, must not be called; if it is, the uncalled mocking helper will alert me by raising an exception.

So the mocking helpers not only temporarily install mock implementations of functions but also assert that those mock implementations are (or are not) called as expected. In the following code, for example, I use this capability to make sure that the mailing-list service is asked to get the subscriber count for the mailing list having the right key. I also simulate the service returning a subscriber count of 123.

def test_unconfirmed_msgs_must_be_confirmed(self):
    """An unconfirmed msg must be confirmed and not sent."""
    mlkey = 'TEST_ML_KEY'
    body = 'Test message!'
    data = self._request_data(mlkey=mlkey, body=body)
    with mocked(mls.subscriber_count, expecting=((mlkey,),{}), returning=123):
        with uncalled(sms.send):
            resp = self.app.request("/mailings", method='POST', data=data)
    self.assertEqual(resp.status, '200 OK')
    self.assertIn('Reply CONFIRM', resp.data)    # we must ask for confirmation
    self.assertIn('123 subscribers', resp.data)  # and supply subscriber count
    self.assertIn(body, resp.data)               # and the msg to be sent

The mocked helper lets you do a few more things, too, but you get the idea: Sometimes a short helper function can take you a long way.