mock_async_func.py

def mock_async_func(func, data):
    """The challenge: Mock an "async" function (ie, a function that takes a
    callback arg that it calls with its result instead of returning it)
    without needing to manually duplicate (and maintain) the original
    function's signature.

    Given this original function:

        def get_page(url, callback):
            libbitly.async_http.http_fetch(url, callback=callback)

    The way we've mocked this kind of function in the past requires duplicated
    the function and its signature in order to get access to the callback
    passed in:

        def make_mock_get_page(data):
            def mock_get_page(url, callback):
                callback(data)
            return mock_get_page

        with mock.patch('get_page', new=make_mock_get_page('page data')):
            pass

    This works, but it's kind of a pain in the ass, especially when mocking
    more than one async function or when the signature of a mocked function
    changes.

    Using this mock_async_func helper, the above turns into:

        with mock.patch('get_page', new=mock_async_func(get_page, 'page data')):
            pass

    There is some redundancy where the thing to be patched must be passed into
    mock.patch as a string and into mock_async_func as its real self. I'm not
    exactly sure what to do about that, though.
    """
    args, varargs, keywords, defaults = inspect.getargspec(func)
    assert 'callback' in args, 'func must take a callback arg'

    # We need to store our mock data in a unique field in the namespace we'll
    # use when executing our generated function to avoid potential name
    # clashes in the function itself.
    data_key = '_data_%s' % str(uuid.uuid4()).replace('-', '_')
    ns = { data_key: data }

    # Now create the code for the generated function, whose name and signature
    # should match the original function but whose body should simply call the
    # callback function with our mock data (which it finds under the unique
    # key created above).
    func_name = func.__name__
    sig = inspect.formatargspec(args, varargs, keywords, defaults)
    src = 'def %s%s: return callback(%s)' % (func_name, sig, data_key)

    # Compile the function into a code object, and exec that code object in
    # our artificial environment to turn it into a real function.
    code = compile(src, '<string>', 'exec')
    exec code in ns

    # Return the new function itself.
    return ns[func_name]