blthayer
blthayer

Reputation: 922

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:

Any help is appreciated, thank you.

Upvotes: 12

Views: 28884

Answers (3)

scharfmn
scharfmn

Reputation: 3661

I think the correct approach is found in this answer

note: the below is a sketch - may not get all the details of the OP correct

import unittest
from unittest.mock import patch
from tmp import my_module

class MyClassTestCase(unittest.TestCase):

    @patch('tmp.my_module.MyClass')
    def test_create_class_call_method(self, my_class_mock):
       my_class_mock.return_value.my_method.return_value = 2
       value = my_module.create_class_call_method()
       self.assertEqual(value, 2)

Upvotes: 5

blthayer
blthayer

Reputation: 922

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

Upvotes: 11

lmiguelvargasf
lmiguelvargasf

Reputation: 70003

I am not sure about the implementation of create_class_call_method, but try the following:

from unittest import mock

class MyClassTestCase(unittest.TestCase):
    @mock.patch('tmp.my_module.MyClass.my_method')
    @mock.patch('tmp.my_module.MyClass.__init__')
    def test_create_class_call_method(self, my_class_init, my_method_mock):
        my_class_init.return_value = None
        my_method.return_value = 2

        value = my_module.create_class_call_method()

        self.assertEqual(value, 2)

Upvotes: 6

Related Questions