Context manager to help monkeypatch global variables used by a function under test. Uses the mock library

gistfile1.py

import types
import mock

def patch_func_globals(func, **vars):
    '''
    Monkeypatches global variables used by a function or method

    It's useful if you need to mock a global function or variable
    your function uses

    Usage:

        with patch_func_globals(your_func, used_func=mocked_func):
            your_func()

    where `used_func` is called inside `your_func()`

    If your function is actually a generator you must call it before
    patching it, like:

        my_generator = my_func_with_yield()
        with patch_func_globals(my_generator, foo=foo_mock):
            result = list(my_generator)  # consume the generator

    Works with functions and methods, even if they where wrapped
    using a decorator (like @memoized)

    '''

    # if it's a generator get the function object from the frame
    if hasattr(func, 'gi_frame'):
        return mock.patch.dict(func.gi_frame.f_globals, vars)

    # methods have the function object on __func__ attribute
    if hasattr(func, '__func__'):
        func = func.__func__

    # if the function is @memoized we need to get the original function
    # from the closure
    func = get_function_from_closure(func, func.__name__) or func

    return mock.patch.dict(func.__globals__, vars)

def get_function_from_closure(func, name):
    '''
    Looks for a function with `name` inside `func`s closure

    If not found returns None

    A closure is a tuple of `cell` objects, the cell object has
    an attribute `cell_contents` which could be any object

    '''
    if not func.__closure__:
        return None

    def is_function_with_name(obj, name):
        return type(obj) == types.FunctionType and \
               obj.__name__ == name

    for cell in func.__closure__:
        try:
            if is_function_with_name(cell.cell_contents, name):
                return cell.cell_contents
        except ValueError: # cell is empty
            pass

    return None