Mock class in Python with decorator patch

Question:

I would like to patch a class in Python in unit testing. The main code is this (mymath.py):

class MyMath:
    def my_add(self, a, b):
        return a + b

def add_three_and_two():
    my_math = MyMath()
    return my_math.my_add(3, 2)

The test class is this:

import unittest
from unittest.mock import patch
import mymath

class TestMyMath(unittest.TestCase):
    @patch('mymath.MyMath')
    def test_add_three_and_two(self, mymath_mock):
        mymath_mock.my_add.return_value = 5

        result = mymath.add_three_and_two()

        mymath_mock.my_add.assert_called_once_with(3, 2)
        self.assertEqual(5, result)

unittest.main()

I am getting the following error:

AssertionError: Expected 'my_add' to be called once. Called 0 times.

The last assert would also fail:

AssertionError: 5 != <MagicMock name='MyMath().my_add()' id='3006283127328'>

I would expect that the above test passes. What I did wrong?

UPDATE:
Restrictions:

  • I would not change the tested part if possible. (I am curious if it is even possible, and this is the point of the question.)
  • If not possible, then I want the least amount of change in the to be tested part. Especially I want to keep the my_add() function non-static.
Asked By: Csaba Faragó

||

Answers:

Your code is almost there, some small changes and you’ll be okay:

  1. my_add should be a class method since self does not really play a role here.
  2. If my_add is an instance method, then it will be harder to trace the calls, since your test will track the instance signature, not the class sig
  3. Since you are are patching, not stubbing, you should use the "real thing", except when mocking the return value.

Here’s what that looks like in your code:

class MyMath:

    @classmethod
    def my_add(cls, a, b):
        return a + b

def add_three_and_two():
    return MyMath.my_add(3, 2)

Now, the test:

import unittest
from unittest.mock import patch, MagicMock
import mymath


class TestMyMath(unittest.TestCase):

    @patch('mymath.MyMath')
    def test_add_three_and_two(self, mymath_mock):

        # Mock what `mymath` would return 
        mymath_mock.my_add.return_value = 5

        # We are patching, not stubbing, so use the real thing
        result = mymath.add_three_and_two()
        mymath.MyMath.my_add.assert_called_once_with(3, 2)
        self.assertEqual(5, result)


unittest.main()

This should now work.

Answered By: Yaakov Bressler

Instead of patching the entire class, just patch the function.

class TestMyMath(unittest.TestCase):
    @patch.object(mymath.MyMath, 'my_add')
    def test_add_three_and_two(self, m):
        m.return_value = 5

        result = mymath.add_three_and_two()

        m.assert_called_once_with(3, 2)
        self.assertEqual(5, result)

I think the original problem is that my_math.my_add produces a new mock object every time it is used; you configured one Mock‘s return_value attribute, but then checked if another Mock instance was called. At the very least, using patch.object ensures you are disturbing your original code as little as possible.

Answered By: chepner