2023/03/21

Mock Open Write On Python

在 python 的測試中,有時候 mock 只要 mock a.b.c 就可以成功,但我們使用了 with 的情況下,就必須借助 __enter__ 這個魔術方法來測試了,我們介紹兩種寫法。

my_open.py

def write(path, context):
    with open(path) as f:
        f.write(context)

test_my_open.py

import unittest
import my_open

from unittest.mock import mock_open, patch


class MyTestCase(unittest.TestCase):
    def test_my_open_1(self):
        with patch('builtins.open', mock_open()) as o:
            path = '/tmp/test.txt'
            context = 'hello'
            my_open.write(path, context)

            o.assert_called_once_with(path)
            o().write.assert_called_once_with(context)

    @patch('builtins.open')
    def test_my_open_2(self, mock_o: mock_open):
        path = '/tmp/test.txt'
        context = 'hello'
        my_open.write(path, context)

        mock_o.assert_called_once_with(path)
        handle = mock_o.return_value.__enter__.return_value
        handle.write.assert_called_once_with(context)


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