Search code examples
pythonpython-3.xunit-testingmockingpython-unittest

Python unittest mock class and class method


I feel like this may be relatively simple, but I'm pulling my hair out to get this working. I'd like to mock an entire class, and then specify the return value for one of this class's methods.

I already looked here, at several other questions, and of course in the docs. I'm still unable to get this to work. Please see my simple example below.

Contents of directory tmp:

tmp
├── __init__.py
├── my_module.py
└── test_my_module.py

Contents of my_module.py:

class MyClass:
    def __init__(self):
        # Do expensive operations that will be mocked in testing.
        self.a = 7

    def my_method(self):
        # For sake of simple example, always return 1.
        return 1


def create_class_call_method():
    """Create MyClass instance and call its my_method method, returning
    the result."""
    instance = MyClass()
    value = instance.my_method()
    return value

Contents of test_my_module.py:

import unittest
from unittest.mock import patch, Mock

from tmp import my_module


class MyClassTestCase(unittest.TestCase):

    def test_create_class_call_method(self):
        # Attempt to patch MyClass as well as specify a return_value for
        # the my_method method (spoiler: this doesn't work)
        with patch('tmp.my_module.MyClass',
                   my_method=Mock(return_value=2)):
            value = my_module.create_class_call_method()

        self.assertEqual(value, 2)


if __name__ == '__main__':
    unittest.main()

Results of running test_my_module.py:

2 != <MagicMock name='MyClass().my_method()' id='140234477124048'>

Expected :<MagicMock name='MyClass().my_method()' id='140234477124048'>
Actual   :2

Some other things I've tried:

  • Rather than ..., my_method=Mock(return_value=2)) in the patch statement, unpack a dictionary like so: **{'my_method.return_value': 2}
  • Nested with patch statements. Outer statement is simple like with patch('tmp.my_module.MyClass'):, inner statement attempts to patch my_method like so: with patch('tmp.my_module.MyClass.my_method, return_value=2)
  • Use patch decorators instead of context managers
  • Change patch statement to with patch('tmp.my_module.MyClass') as p: and then inside the with statement, try to set p like so: p.evaluate = Mock(return_value=2)

Any help is appreciated, thank you.


Solution

  • I've found a much better solution. In short, we need to mock out the return_value of the MyClass mock. Here's the working test code:

    import unittest
    from unittest.mock import patch, Mock, MagicMock
    
    from tmp import my_module
    
    
    class MyClassTestCase(unittest.TestCase):
    
        def test_create_class_call_method(self):
            # Create a mock to return for MyClass.
            m = MagicMock()
            # Patch my_method's return value.
            m.my_method = Mock(return_value=2)
    
            # Patch MyClass. Here, we could use autospec=True for more
            # complex classes.
            with patch('tmp.my_module.MyClass', return_value=m) as p:
                value = my_module.create_class_call_method()
    
            # Method should be called once.
            p.assert_called_once()
            # In the original my_method, we would get a return value of 1.
            # However, if we successfully patched it, we'll get a return
            # value of 2.
            self.assertEqual(value, 2)
    
    
    if __name__ == '__main__':
        unittest.main()
    

    And the successful results:

    Ran 1 test in 0.002s
    
    OK