Any way to reset a mocked method to its original state? - Python Mock - mock 1.0b1

Question

I have the following simplified class I'm mocking:

class myClass(object):
    @staticmethod
    def A():
        #...

    def check(self):
        #code...
        value = self.A()
        #more code...

In my first test I mock only the method A

from django.test import TestCase
from mock import MagicMock
import myClass

class FirstTest(TestCase):

def setUp(self):
    myClass.A = MagicMock(return_value = 'CPU')

def test(self):
    #some tests 
    myClassObj = myClass()
    myClassObj.check()

Whereas in my second test I mock the entire check method:

from django.test import TestCase
from mock import MagicMock
import myClass

class SecondTest(TestCase):

def setUp(self):
    myClass.check = MagicMock(return_value = someObject)

def test(self):
    #some tests 
    myClassObj = myClass()
    myClassObj.check()

Now my assertions from my first test fail because, instead of calling check() and mocking A() inside check(), it calls the completely mocked check() from my second test.

Is there any way to clear and set the method to be 'normal' after the test? I tried myClass.check.reset_mock() already, but it doesn't seem to do anything. Moving the order of my tests doesn't do anything either.

I'm using mock 1.0b1 for python from http://pypi.python.org/pypi/mock/

Answer 1

You can stash the function away on self and put it back when you're done.

import unittest

from mock import MagicMock
from MyClass import MyClass

class FirstTest(unittest.TestCase):

    def setUp(self):
        self.A = MyClass.A
        MyClass.A = MagicMock(name='mocked A', return_value='CPU')


    def tearDown(self):
        MyClass.A = self.A

    def test_mocked_static_method(self):
        print 'First Test'
        print MyClass.check
        print MyClass.A


class SecondTest(unittest.TestCase):

    def setUp(self):
        MyClass.check = MagicMock(name='mocked check', return_value=object)

    def test_check_mocked_check_method(self):
        print 'Second Test'
        print MyClass.check
        print MyClass.A


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

Running this file gives the following output:

First Test
<unbound method MyClass.check> 
<MagicMock name='mocked A' id='141382732'>
Second Test
<MagicMock name='mocked check' id='141382860'>
<unbound method MyClass.A>

I found myself using the patch decorator a lot more than setUp and tearDown now. In this case you could do

from mock import patch

@patch('MyClass.A')
def test_mocked_static_method(self, mocked_A)
    mocked_A.return_value = 'CPU'
    # This mock will expire when the test method is finished

Answer 2

You can use mock.patch as a decorator or a context manager:

from mock import patch, MagicMock

@patch('myClass.A', MagicMock(return_value='CPU'))
def test(self):
    pass

or:

def test(self):
    with patch('myClass.A', MagicMock(return_value='CPU')):
        pass

If you don't supply a mock object to patch then it will provide an autospecced mock that you can modify:

@patch('myClass.A')
def test(self, mock_A):
    mock_A.return_value = 'CPU'
    pass

or:

def test(self):
    with patch('myClass.A') as mock_A:
        mock_A.return_value = 'CPU'
        pass

In all cases the original value will be restored when the decorated test function or context manager finishes.