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

Solution 1:

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.

Solution 2:

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