How to mock and test decorator?

Question

How to test the following decorator which calls 3rd party library?

import third_party_lib
import functools

class MyDecorator:
    def __init__(self, ....):
        self.third_party_lib = ThirdPartyLib(....) # will create a 3rd party instance

    def __call__(self, ...):
        def decorator(f):
            @functools.wraps(f)
            def wrap(*arg, **kwargs):
                result = f(*arg, **kwargs)
                # ....
                a = self.third_party_lib.send(value=result).get()
                # ....
                return result
            return wrap
        return decorator

I need to create an unit test to assert third_party_lib.send() is called if a function is decorated by the decorator. And ideally, also assure the result of the test function is passed to the function.

decorator = MyDecorator(....)

@decorator(....)
def test_func():
    ret = ...
    return ret # ret should be passed to `third_party_lib.send()`

Answer 1

If you want to verify that the thirdparty function is called correctly, you can mock it and check that the mock is called with the correct parameters. As the ThirdPartyLib intialization shall also be mocked, as mentioned in the comments, you hav to make sure that the docorator is constructed after the mock has been set, for example by constructing it inside the test:

from unittest import mock

@mock.patch('third_party_lib.ThirdPartyLib')
def test_my_decorator(mocked_lib):
    decorator = MyDecorator()

    @decorator()
    def example_func():
        return 42

    example_func()
    mocked_lib.return_value.send.assert_called_once_with(value=42)

If you need the decorated function in more tests, you can wrap it in a function instead:

def decorated_func():
    decorator = MyDecorator()

    @decorator()
    def example_func():
        return 42

    return example_func

@mock.patch('third_party_lib.ThirdPartyLib')
def test_my_decorator(mocked_lib):
    decorated_func()()
    mocked_lib.return_value.send.assert_called_once_with(value=42)