Search code examples
pythonpython-unittestpython-mockpython-unittest.mock

Mock a subset of a Python class's methods and properties


I'm using the mock Python module for performing my tests.

There are times when I'm mocking a class, however I just want to mock some of its methods and properties, and not all of them.

Suppose the following scenario:

# module.py
class SomeClass:
    def some_method(self):
        return 100

    def another_method(self):
        return 500

# test.py
class Tests(unittest.TestCase):
    @patch('module.SomeClass')
    def test_some_operation(self, some_class_mock):
        some_class_instance = some_class_mock.return_value

        # I'm mocking only the some_method method.
        some_class_instance.some_method.return_value = 25

        # This is ok, the specific method I mocked returns the value I wished.
        self.assertEquals(
            25,
            SomeClass().some_method()
        )

        # However, another_method, which I didn't mock, returns a MagicMock instance
        # instead of the original value 500
        self.assertEquals(
            500,
            SomeClass().another_method()
        )

On the code above, once I patch the SomeClass class, calls to methods whose return_values I didn't exlicitely set will return MagicMock objects.

My question is: How can I mock only some of a class methods but keep others intact?

There are two ways I can think of, but none of them are really good.

  1. One way is to set the mock's method to the original class method, like this:

    some_class_instance.another_method = SomeClass.another_method
    

    This is not really desirable because the class may have a lot of methods and properties to "unmock".

  2. Another way is to patch each method I want explicitly, such as:

     @patch('module.SomeClass.some_method')
     def test_some_operation(self, some_method_mock):
    

    But this doesn't really work if I want to mock the class itself, for mocking calls to the initializer for example. The code below would override all SomeClass's methods anyway.

     @patch('module.SomeClass.some_method')
     @patch('module.SomeClass')
     def test_some_operation(self, some_class_mock, some_method_mock):
    

Here is a more specific example:

class Order:
    def process_event(self, event, data):
        if event == 'event_a':
            return self.process_event_a(data)

        elif event == 'event_b':
            return self.process_event_b(data)

        else:
            return None

    def process_event_a(self, data):
        # do something with data

    def process_event_b(self, data):
        # do something different with data

In this case, I have a general method process_event which calls a specific processing event depending on the supplied event.

I would like to test only the method process_event. I just want to know if the proper specific event is called depending on the event I supply.

So, in my test case what I want to do is to mock just process_event_a and process_event_b, call the original process_event with specific parameters, and then assert either process_event_a or process_event_b were called with the proper parameters.


Solution

  • Instead of patching the whole class, you must patch the object. Namely, make an instance of your class, then, patch the methods of that instance.

    Note that you can also use the decorator @patch.object instead of my approach.

    class SomeClass:
        def some_method(self):
            return 100
    
        def another_method(self):
            return 500
    

    In your test.py

    from unittest import mock
    
    class Tests(unittest.TestCase):
        
    
        def test_some_operation(self):
            some_class_instance = SomeClass()
    
            # I'm mocking only the some_method method.
            with mock.patch.object(some_class_instance, 'some_method', return_value=25) as cm:
    
                # This is gonna be ok
                self.assertEquals(
                    25,
                    SomeClass().some_method()
                )
    
                # The other methods work as they were supposed to.
                self.assertEquals(
                    500,
                    SomeClass().another_method()
                )