PicklingError when using mock to check function is called

user_8576714 0
Viewed 3851

I'm having trouble testing that a method is called using mock -- as a simple example, let's say that the method is os.getcwd. I want to test that my own function, pickle_wdir, is calling os.getcwd as intended. However, the function I am testing pickles the value returned by os.getcwd, which results in an error.

Here's a simple example to reproduce the error.

os_ex.py:

import os
import pickle

def pickle_wdir(filename):
    dir = os.getcwd()
    with open(filename, 'wb') as handle:
        pickle.dump(dir, handle)

test_os_ex.py:

from unittest import TestCase
from unittest.mock import patch
from os_ex import pickle_wdir


class TestPickleWdir(TestCase):
    def test_os_called(self):
        fname = 'dir.pickle'
        with patch('os_ex.os') as mocked_obj:
            pickle_wdir(fname)
            mocked_obj.getcwd.assert_called()

The error message returned is

_pickle.PicklingError: Can't pickle <class 'unittest.mock.MagicMock'>: it's not the same object as unittest.mock.MagicMock.

How can I test that os.getcwd is called without getting this PicklingError?

python-unittest.mock python-unittest pickle python
1 Answers

You need to mock os.getcwd(), open() and pickle.dump() methods. You can use unittest.mock.patch as context manager to do this.

E.g.

os_ex.py:

import os
import pickle


def pickle_wdir(filename):
    dir = os.getcwd()
    with open(filename, 'wb') as handle:
        pickle.dump(dir, handle)

test_os_ex.py:

import unittest
from unittest.mock import patch, mock_open
from os_ex import pickle_wdir


class TestOsEx(unittest.TestCase):
    def test_pickle_wdir(self):
        fname = 'dir.pickle'
        m = mock_open(read_data='mocked data')
        with patch('os_ex.os') as mocked_obj, patch('builtins.open', m) as mocked_open, patch('pickle.dump') as mocked_dump:
            mocked_obj.getcwd.return_value = '/root'
            pickle_wdir(fname)
            mocked_obj.getcwd.assert_called()
            m.assert_called_with(fname, 'wb')
            handle = mocked_open()
            mocked_dump.assert_called_with('/root', handle)


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

unit test results with 100% coverage:

.
----------------------------------------------------------------------
Ran 1 test in 0.012s

OK
Name                                       Stmts   Miss  Cover   Missing
------------------------------------------------------------------------
src/stackoverflow/60627827/os_ex.py            6      0   100%
src/stackoverflow/60627827/test_os_ex.py      17      0   100%
------------------------------------------------------------------------
TOTAL                                         23      0   100%
user_64635580
Related