CODEWITHSUNDEEP
pythonpython-unittest
mojoken
mojoken

Reputation: 1362

How to write a mock that works with decorators?

I have existing unit tests which mock methods of a python class, and they have worked successfully for some time. Now I am adding an AWS xray decorator onto one of my mocked methods and the calls to this mock are failing with "missing a required argument" exception.

Here is part of the class with the myMethod() function which takes 2 arguments. The xray_recorder.capture decorator before the function definition is what's newly added:

@xray_recorder.capture()    
def myMethod(self, dbConn, userSession):
    ...

In my unit test, I have mocked this class and function like follows:

mock_myClass = mock.create_autospec(myClass)
mock_myClass.myMethod.return_value = {
    'attr1': 'value1',
    ...
}

When the caller invokes the method in regular usage, it works fine. But in my unit test it fails. Here is the call to the method:

            retVal = mock_myClass.myMethod(dbConn, userSession)

This call now raises an exception: missing a required argument: 'userSession'

I also tried defining the mock myMethod with args and kwargs:

def mocked_myMethod(*args, **kwargs): return { 'attr1': 'value1', ... }

and using:

mock_myClass.myMethod.side_effect = mocked_myMethod

but this fails in the same way.

My understanding is that decorators have to pass the argument list through to the function they're invoking. If I look at the source code for xray_recorder.capture() from the aws_xray_sdk I see:

def capture(self, name=None):
    """
    A decorator that records enclosed function in a subsegment.
    It only works with synchronous functions.

    params str name: The name of the subsegment. If not specified
    the function name will be used.
    """
    return self.in_subsegment(name=name)

...

def in_subsegment(self, name=None, **subsegment_kwargs):
    """
    Return a subsegment context manger.

    :param str name: the name of the subsegment
    :param dict subsegment_kwargs: remaining arguments passed directly to `begin_subsegment`
    """
    return SubsegmentContextManager(self, name=name, **subsegment_kwargs)

How can I make my mock work with this?

Upvotes: 1

Views: 294

Answers (1)

Diddi Oskarsson
Diddi Oskarsson

Reputation: 78

You need to call myMethod on the mocked object, not the class itself for it to work.

mock_myClass = mock.create_autospec(myClass)
...
retVal = mock_myClass.myMethod(dbConn, userSession) # Note use of mock_MyClass

Calling myClass.myMethod(...) would call the method on the class directly, bypassing the mock.

Upvotes: 1

Related Questions