雑念ストレージ

プログラミング関連のメモとか

【Python】mock_openで複数のファイルの読み書きをMock化する

前にmock_openを使って1つのファイルの読み書きをMock化する方法を書いた。
複数のファイルの読み書きをMock化する場合は少しコツがあったので、メモしておく。

複数ファイルの読み込みをMock化する場合

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


def read_2files(file_path1: str, file_path2: str) -> tuple[str, str]:
    # ファイルを2つ読み込んで返すだけ
    with open(file_path1, 'r', encoding='UTF-8') as f:
        data1 = f.read()

    with open(file_path2, 'r', encoding='UTF-8') as f:
        data2 = f.read()

    return data1, data2


class TestRead2Files(unittest.TestCase):
    def test_read_2files(self):
        # ファイルのモックを準備する
        file_dict = {
            'file1.txt': mock_open(read_data='AAA')(),
            'file2.txt': mock_open(read_data='BBB')(),
        }

        mock = MagicMock()

        # args[0]にファイルパスが入っているので、それに合わせてファイルのモックを切り替える。
        # openを呼び出す際に名前付き引数を使っているなら、lambdaの引数はargsだけではなくkwargsも用意する必要があるので注意。
        # 今回の場合、encodingを名前付き引数で指定している。
        mock.side_effect = lambda *args, **kwargs: file_dict[args[0]]

        with patch('builtins.open', mock):
            # 処理を呼び出す
            result = read_2files('file1.txt', 'file2.txt')
            # 結果を確認
            self.assertEqual(result, ('AAA', 'BBB'))

以下の部分で、ファイルパスごとにread_dataを設定したファイルのモックを用意している。

file_dict = {
    'file1.txt': mock_open(read_data='AAA')(),
    'file2.txt': mock_open(read_data='BBB')(),
}

mock_open(read_data='AAA')()のところは、カッコが続いて面食らうけど、以下のように段階を踏んで考えるとわかりやすい。

  • mock_open(read_data='AAA')
    • open関数のモック
  • mock_open(read_data='AAA')()
    • open関数を呼び出した結果
    • つまり、with open(file_path1, 'r', encoding='UTF-8') as f: でいう f のこと

複数ファイルの書き込みをMock化する場合

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


def write_2files(file_path1: str, file_path2: str):
    # ファイル2つに書き込むだけ
    with open(file_path1, 'w', encoding='UTF-8') as f:
        f.write("AAA")

    with open(file_path2, 'w', encoding='UTF-8') as f:
        f.write("BBB")


class TestWrite2Files(unittest.TestCase):

    def test_write_multi(self):
        file_mock1 = mock_open()()
        file_mock2 = mock_open()()

        file_dict = {
            'file1.txt': file_mock1,
            'file2.txt': file_mock2,
        }

        mock = MagicMock()
        mock.side_effect = lambda *args, **kwargs: file_dict[args[0]]

        with patch('builtins.open', mock):
            write_2files('file1.txt', 'file2.txt')

            file_mock1.write.assert_called_once_with("AAA")
            file_mock2.write.assert_called_once_with("BBB")

読み込みとあまり変わらない。
書き込み(write)が呼ばれたことをassertしたいので、ファイルのモックは変数に持たせている。(file_mock1file_mock2)

複数ファイルの読み書き両方をMock化する場合

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


def read_write_multi(input_file_path1: str, input_file_path2: str,
                     output_file_path1: str, output_file_path2: str):
    # ファイル2つを読み込んで、それぞれの内容に追記して別ファイル2つに書き出す
    with open(input_file_path1, 'r', encoding='UTF-8') as f:
        data1 = f.read()

    with open(input_file_path2, 'r', encoding='UTF-8') as f:
        data2 = f.read()

    with open(output_file_path1, 'w', encoding='UTF-8') as f:
        f.write(data1 + " APPEND 1")

    with open(output_file_path2, 'w', encoding='UTF-8') as f:
        f.write(data2 + " APPEND 2")


class TestReadWriteMulti(unittest.TestCase):

    def test_read_write_multi(self):
        output_file_mock1 = mock_open()()
        output_file_mock2 = mock_open()()

        file_dict = {
            'input1.txt': mock_open(read_data='AAA')(),
            'input2.txt': mock_open(read_data='BBB')(),
            'output1.txt': output_file_mock1,
            'output2.txt': output_file_mock2,
        }

        mock = MagicMock()
        mock.side_effect = lambda *args, **kwargs: file_dict[args[0]]

        with patch('builtins.open', mock):
            read_write_multi('input1.txt', 'input2.txt', 'output1.txt', 'output2.txt')

            output_file_mock1.write.assert_called_once_with("AAA APPEND 1")
            output_file_mock2.write.assert_called_once_with("BBB APPEND 2")

読み込みの例と書き込みの例を組み合わせただけ。