Can I patch a Python decorator before it wraps a function?

54,158

Solution 1

Decorators are applied at function definition time. For most functions, this is when the module is loaded. (Functions that are defined in other functions have the decorator applied each time the enclosing function is called.)

So if you want to monkey-patch a decorator, what you need to do is:

  1. Import the module that contains it
  2. Define the mock decorator function
  3. Set e.g. module.decorator = mymockdecorator
  4. Import the module(s) that use the decorator, or use it in your own module

If the module that contains the decorator also contains functions that use it, those are already decorated by the time you can see them, and you're probably S.O.L.

Edit to reflect changes to Python since I originally wrote this: If the decorator uses functools.wraps() and the version of Python is new enough, you may be able to dig out the original function using the __wrapped__ attribute and re-decorate it, but this is by no means guaranteed, and the decorator you want to replace also may not be the only decorator applied.

Solution 2

It should be noted that several of the answers here will patch the decorator for the entire test session rather than a single test instance; which may be undesirable. Here's how to patch a decorator that only persists through a single test.

Our unit to be tested with the undesired decorator:

# app/uut.py

from app.decorators import func_decor

@func_decor
def unit_to_be_tested():
    # Do stuff
    pass

From decorators module:

# app/decorators.py

def func_decor(func):
    def inner(*args, **kwargs):
        print "Do stuff we don't want in our test"
        return func(*args, **kwargs)
    return inner

By the time our test gets collected during a test run, the undesired decorator has already been applied to our unit under test (because that happens at import time). In order to get rid of that, we'll need to manually replace the decorator in the decorator's module and then re-import the module containing our UUT.

Our test module:

#  test_uut.py

from unittest import TestCase
from app import uut  # Module with our thing to test
from app import decorators  # Module with the decorator we need to replace
import imp  # Library to help us reload our UUT module
from mock import patch


class TestUUT(TestCase):
    def setUp(self):
        # Do cleanup first so it is ready if an exception is raised
        def kill_patches():  # Create a cleanup callback that undoes our patches
            patch.stopall()  # Stops all patches started with start()
            imp.reload(uut)  # Reload our UUT module which restores the original decorator
        self.addCleanup(kill_patches)  # We want to make sure this is run so we do this in addCleanup instead of tearDown

        # Now patch the decorator where the decorator is being imported from
        patch('app.decorators.func_decor', lambda x: x).start()  # The lambda makes our decorator into a pass-thru. Also, don't forget to call start()          
        # HINT: if you're patching a decor with params use something like:
        # lambda *x, **y: lambda f: f
        imp.reload(uut)  # Reloads the uut.py module which applies our patched decorator

The cleanup callback, kill_patches, restores the original decorator and re-applies it to the unit we were testing. This way, our patch only persists through a single test rather than the entire session -- which is exactly how any other patch should behave. Also, since the clean up calls patch.stopall(), we can start any other patches in the setUp() we need and they will get cleaned up all in one place.

The important thing to understand about this method is how the reloading will affect things. If a module takes too long or has logic that runs on import, you may just need to shrug and test the decorator as part of the unit. :( Hopefully your code is better written than that. Right?

If one doesn't care if the patch is applied to the whole test session, the easiest way to do that is right at the top of the test file:

# test_uut.py

from mock import patch
patch('app.decorators.func_decor', lambda x: x).start()  # MUST BE BEFORE THE UUT GETS IMPORTED ANYWHERE!

from app import uut

Make sure to patch the file with the decorator rather than the local scope of the UUT and to start the patch before importing the unit with the decorator.

Interestingly, even if the patch is stopped, all the files that already imported will still have the patch applied to the decorator, which is the reverse of the situation we started with. Be aware that this method will patch any other files in the test run that are imported afterwards -- even if they don't declare a patch themselves.

Solution 3

When I first ran across this problem, I use to rack my brain for hours. I found a much easier way to handle this.

This will fully bypass the decorator, like the target wasn't even decorated in the first place.

This is broken down into two parts. I suggest reading the following article.

http://alexmarandon.com/articles/python_mock_gotchas/

Two Gotchas that I kept running into:

1.) Mock the Decorator before the import of your function/module.

The decorators and functions are defined at the time the module is loaded. If you do not mock before import, it will disregard the mock. After load, you have to do a weird mock.patch.object, which gets even more frustrating.

2.) Make sure you are mocking the correct path to the decorator.

Remember that the patch of the decorator you are mocking is based on how your module loads the decorator, not how your test loads the decorator. This is why I suggest always using full paths for imports. This makes things a lot easier for testing.

Steps:

1.) The Mock function:

from functools import wraps

def mock_decorator(*args, **kwargs):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            return f(*args, **kwargs)
        return decorated_function
    return decorator

2.) Mocking the decorator:

2a.) Path inside with.

with mock.patch('path.to.my.decorator', mock_decorator):
     from mymodule import myfunction

2b.) Patch at top of file, or in TestCase.setUp

mock.patch('path.to.my.decorator', mock_decorator).start()

Either of these ways will allow you to import your function at anytime within the TestCase or its method/test cases.

from mymodule import myfunction

2.) Use a separate function as a side effect of the mock.patch.

Now you can use mock_decorator for each decorator you want to mock. You will have to mock each decorator separately, so watch out for the ones you miss.

Solution 4

To patch a decorator, you need to either import or reload the module which uses that decorator after patching it OR redefine the module's reference to that decorator altogether.

Decorators are applied at the time that a module is imported. This is why if you imported a module which uses a decorator you want to patch at the top of your file and attempt it to patch it later on without reloading it, the patch would have no effect.

Here is an example of the first way mentioned of doing this - reloading a module after patching a decorator it uses:

import moduleA
...

  # 1. patch the decorator
  @patch('decoratorWhichIsUsedInModuleA', examplePatchValue)
  def setUp(self)
    # 2. reload the module which uses the decorator
    reload(moduleA)

  def testFunctionA(self):
    # 3. tests...
    assert(moduleA.functionA()...

Helpful References:

Solution 5

The following worked for me:

  1. Eliminate the import statement that loads the test target.
  2. Patch the decorator on test startup as applied above.
  3. Invoke importlib.import_module() immediately after patching to load the test target.
  4. Run tests normally.

It worked like a charm.

Share:
54,158
Chris Sears
Author by

Chris Sears

Updated on July 10, 2022

Comments

  • Chris Sears
    Chris Sears almost 2 years

    I have a function with a decorator that I'm trying test with the help of the Python Mock library. I'd like to use mock.patch to replace the real decorator with a mock 'bypass' decorator which just calls the function.

    What I can't figure out is how to apply the patch before the real decorator wraps the function. I've tried a few different variations on the patch target and reordering the patch and import statements, but without success. Any ideas?

  • Paragon
    Paragon about 11 years
    The following wasted quite a bit of my time: keep in mind that Python only imports modules once. If you're running a suite of tests, trying to mock a decorator in one of your tests, and the decorated function is imported elsewhere, mocking the decorator will have no effect.
  • IxDay
    IxDay over 9 years
    use the builtin reload function in order to regenerate the python binary code docs.python.org/2/library/functions.html#reload and monkeypatch your decorator
  • Michele d'Amico
    Michele d'Amico over 8 years
    I see a lot issues in this answer. The first (and the bigger one) is that you cannot have access to the original function if it is decorated yet (that is the OP issue). Moreover you don't remove the patch after test is done and that can cause problems when you run it in a test suite.
  • Malcolm Jones
    Malcolm Jones over 7 years
    user2859458, this helped me out significantly. Accepted answer is good, but this spelled things out for me in a meaningful manner, and included multiple use cases where you might want something slightly different.
  • claytond
    claytond over 6 years
    Ran into the issue reported by @Paragon and worked around it by patching my decorator in the test directory's __init__. That ensured that the patch was loaded before any test file. We have an isolated tests folder so the strategy works for us, but this may not work for every folder layout.
  • Geekfish
    Geekfish over 6 years
    Thank you for this response! In case this is useful to others, I made an extension of patch that will still work as a context manager and do the reloading for you: gist.github.com/Geekfish/aa43368ceade131b8ed9c822d2163373
  • ritratt
    ritratt over 4 years
    After reading this several times, I am still confused. This needs a code example!
  • ritratt
    ritratt over 4 years
    The blog post you cited helped me understand this much better!
  • Srivathsa
    Srivathsa over 4 years
    @claytond Thanks your solution worked for me as I had an isolated tests folder!
  • Dougyfresh
    Dougyfresh over 4 years
    Using the wrapped attribute was the way to go for me, thanks!. I use while hasattr(target_func,'__wrapped__'): target_func=target_func.__wrapped__ to dig out the core function and test it.
  • Michael L Rhyndress
    Michael L Rhyndress over 3 years
    This solved my problem, thanks!! lambda *x, **y: lambda f: f
  • Luis Blanche
    Luis Blanche about 3 years
    The __wrapped__ , i think it is the most elegant way to do this and also the most 'local' way, meaning you just choose not to apply it a the function level
  • mike rodent
    mike rodent over 2 years
    I managed to do something along these lines using importlib.reload (of the module which uses the decorator, after having injected a mock decorator). But don't you then need to have a teardown function which restores the status quo ante? Aren't you otherwise stuck with a mock decorator for all eternity (or at least until the present test suite finishes)? A true mock/patch would end when the test ends, wouldn't it?
  • Nairum
    Nairum over 2 years
    Tests and testing must not affect the production code. In the production code must be only the required code. It is not a nice design. Mocking must be used instead.