'Python patch internal class __call__ method

I have a class B that I want to unit test. The class internally instantiates another class A and calls it via its call method. How do I patch class A and its init and __call__methods?

class A:
    def __init__():
        ....
        
    def __call__():
        ....

class B:
    
    def my_method():
        a = A()
        a()
        
def test_b_my_method():
...


Solution 1:[1]

You can use decorators for that


class A:
    def __init__(self):
        self._foo = 1

    def __call__(self):
        return self._foo + 1

class B:
    def my_method(self):
        a = A()
        return a()

def patch_init(fn):
    def wrap(*args, **kwargs):
        self = args[0]
        fn(*args, **kwargs) # calls A.__init__(*args, **kwargs)
        self._foo = 2
    return wrap

def patch_call(fn):
    def wrap(*args, **kwargs):
        res = fn(*args, **kwargs) # calls A.__call__(*args, **kwargs)
        return res * 2
    return wrap

if __name__ == "__main__":
    # unpatched
    b = B()
    print(b.my_method()) # prints 2 (self._foo + 1)
    # patched
    A.__init__ = patch_init(A.__init__)
    A.__call__ = patch_call(A.__call__)
    b = B()
    print(b.my_method()) # prints 6 (self._foo + 1) * 2

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1