Python mock: special-cases

combining mocks into single decorator

Sometimes mocks (such as a context-manager) have a lot of boilerplate code. To combine several related mocks into one decorator, you can create a custom decorator. This has several advantages:

• do not need arguments for all mocks on all unittests (you choose what to display)
• can provide more meaningful arguments in place of the mocks for other things you actually want to modify

Check this out:

def database_mock( method ):
    conn        = mock.MagicMock()
    cursor      = mock.MagicMock()
    conn.cursor = mock.Mock( return_value=cursor )

    @mock.patch.object( Connection, '__init__',   return_value=None )
    @mock.patch.object( Connection, '__exit__',   return_value=None )
    @mock.patch.object( Connection, '__enter__',  return_value=conn )
    @functools.wraps( method )                                            ## important- so test still referred to as test_* instead of _patched_method
    def _patched_method( self, _enter, _exit, _init, *args, **kwds ):
        return method( self, conn, cursor, *args,**kwds)

    return _patched_method

and it is used like:

class TestSomething( unittest.TestCase ):
    @database_mock
    def test_something(self, conn, cursor):
        cursor.execute.side_effect = (1,1,2,3,0)
        run_something()

mock with statements (context managers)

mock object

cursor = mock.Mock()
db = mock.MagicMock()  # must be MagicMock!
db.dict_cursor.return_value.__enter__.return_value = cursor
with db.dict_cursor() as dcursor:
    assert dcursor == cursor

mock.patch

def myfunc():
    with tempfile.NamedTemporaryFile(prefix='fileprefix') as fh:
        return fh.name


def test_myfunc(mocker):
    mocker.patch('tempfile.NamedTemporaryFile').return_value.__enter__.return_value.name = 'tempfilename'
    assert myfunc() == 'tempfilename'

old method

In order to mock a class used in a with statement, you need to mock it's magic methods (__enter__,__exit__). Here is some code, and the way that it would be mocked.

#### ../my/module.py
from dbinterface import Connection

with Connection( profile='mayadb', database='general' ) as conn:
	cursor = conn.cursor()
	...

conn   = mock.MagicMock()
cursor = mock.Mock()
conn.cursor.return_value    = cursor
cursor.execute.side_effect  = (1,1,0)
cursor.fetchall.side_effect = ([[2]],[[1]])

with mock.patch.object( Connection, '__enter__', return_value=conn ):
    with mock.patch.object( Connection, '__exit__', return_value=None ):
        with mock.patch.object( Connection, '__init__', return_value=None ):
            # ... your test ...

yield different return values over repeated calls

method 1 pass iterable to side_effect

testobj = mock.Mock()
testobj.side_effect = [ 0, 1, 2 ]

testobj() # 0
testobj() # 1
testobj() # 2

method 2 This is much simpler than it might at first seem, simply replace the function with an iterator.

class Test_MyTest( unittest.TestCase ):
	def test_something(self):

		## Results you want to Return
		##
		next_results = [
				'first  iteration',
				'second iteration',
				'third  iteration',
			]


		## Function that returns results
		##
		def yield_results( *args ):
			next_results.pop(0)


		## Mock your class 'MyTest's method 'yield_results'
		##
		with mock.patch.object( MyTest, 'yield_results' ) as _mock_yield:
			_mock_yield.side_effect = yield_results

			MyTest()

		
		## while running MyTest, each time yield_results 
		## is called it will move on to the next result.

datetime

datetime is tricky to mock, since it is written in C.

from datetime import datetime

mock_dt = mock.Mock(side_effect=datetime)
mock_dt.now = mock.Mock(return_value=datetime(1970,1,1,0,0,0))

with mock.patch('{}.datetime.datetime'.format(mymodule.__name__), mock_dt):
    mymodule.test_function()

If you are using isinstance this gets a bit more tricky:

NOTE:

This may be outdated

class Test_MyTest( unittest.TestCase ):
	def test_something(self):
		from datetime import datetime
		now = datetime.now()

		with mock.patch.object( datetime, 'datetime', 
											mock.Mock( wraps=datetime.datetime, return_value=now  )
									) as _mock_datetime:

class MyObj(object):
	def _datetime_now(self):
		from datetime import datetime
		return datetime.now()

open

mock read

import mymodule # module to test

with mock.patch( '{}.open'.format( mymodule.__name__ ), mock.mock_open(read_data-='aaa')) as m:
    mymodule.function() # open will return 'aaa' as it's read-data

mock write

m = mock.mock_open()
with patch('__main__.open', m, create=True):
    with open('foo', 'w') as fd:
        fd.write('some stuff')

m.assert_called_with('foo', 'w')             # m is open()
m().write.assert_called_with('some_stuff')   # m() is the file-descriptor

as decorator

@mock.patch('bultins.open', new_callable=mock.mock_open)
def test_read(self, m_open):
    m_open.read_data = b'abcdefg'

os.environ

with mock.patch.dict('os.environ', {'test': '1'}):
    # your test

with mock.patch('os.environ', new_callable=mock.PropertyMock(return_value={'TEST': 'TRUE'})):
    # your test

replace module

replace os.path with ntpath.

import ntpath
with mock.patch(yourmod.__name__ + '.os.path', new_callable=mock.PropertyMock(return_value=ntpath)):
    yourmod.format_path('/a/b/c')

mock class attribute

@mock.patch('module.Class.attribute', new_callable=mock.PropertyMock)
def test_attribute(self, m_attribute):
    m_attribute.return_value = 'abc'

os.walk

import contextlib

@contextlib.contextmanager
def mock_walk(walk_paths=None):
    walk_paths = walk_paths or {}
    def walk_results(srcpath, *args, **kwargs):
        if srcpath not in walk_paths:
            return os.walk(srcpath, *args, **kwargs)
        for copyfile in walk_paths[srcpath]:
            yield copyfile

    with mock.patch('os.walk', side_effect=walk_results) as mock_walk:
        yield mock_walk


class TestSomething:
    def test_something(self):
        walk_paths = {
            '/dst': [('/dst', ['/a'], ['x0.txt']),
                     ('/dst/a', [], ['x1.txt'])],
        }
        with mock_walk(walk_paths):
            # ...