Search code examples
pythonunit-testingmockingbuilt-in

Mocking selective file writes in python unittest


I looked around a bit on SO but didn't find what I am looking for, I am pretty sure this has been answered elsewhere. So I have two file writes in my function like this:

def write_files():
    with open("a.txt", 'w') as f_h:
      f_h.write("data1")
    with open("b.txt", 'w') as f_h:
      f_h.write("data2")

How do I mock the f_h.write() selectively so that one returns an exception and the other does not? I tried to set side_effect but its not clear where that fits. The test code that I have experimented with has something like this:

from unittest.mock import patch, call, mock_open
import unittest

class Tester(unittest.TestCase):
    def test_analyze(self):
        with patch("builtins.open", mock_open(read_data="data")) as mf:
           # mf.side_effect = [None, Exception()] ?
           write_files()

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

Solution

  • Two things: you have to mock the context manager, e.g. the result of __enter__, and you have to put the side effect on the write method of the mocked file handle (e.g. the result of the __enter__ call):

    class Tester(unittest.TestCase):
        def test_analyze(self):
            with patch("builtins.open", mock_open(read_data="data")) as mf:
                fh_mock = mf.return_value.__enter__.return_value
                fh_mock.write.side_effect = [None, Exception]
    
                with self.assertRaises(Exception):
                    write_files()