unittest Mockの使い方をメモ

背景

pythonのunittest.Mockの使い方をまとめておく

実装

mockを引数経由で渡せる場合

対象の実装

domain/hoge.py

# obj.execute() が例外投げる場合、retry_countの回数だけ再実行するメソッド

class Hoge:
    def retry_with_error(self, obj, retry_count: int = 2) -> None:
        count = 0

        while True:
            try:
                obj.execute()
                return
            except Exception as e:
                count += 1
                if count > retry_count:
                    raise e

テストの実装

tests/hoge_test.py

from domain import Hoge
import unittest
from unittest.mock import MagicMock


class TestHoge(unittest.TestCase):

    def test_retry_with_error(self):

        # execute()が1回だけ呼ばれることを確認
        mock = MagicMock()
        Hoge().retry_with_error(mock)
        mock.execute.assert_called_with() # 実行されることを確認

        # 初回のexecuteでException発生する場合は、retryする
        # 2回目のexecuteは例外を投げずに終了する
        # side_effectを使って動作を制御する
        mock2 = MagicMock()
        mock2.execute = MagicMock(side_effect=[Exception('first error'), None])
        Hoge().retry_with_error(mock2)
        self.assertEqual(mock2.execute.call_count, 2)

        # executeで毎回Exception発生する場合は、3回実行して例外投げる
        with self.assertRaises(Exception):
            mock3 = MagicMock()
            mock3.execute = MagicMock(side_effect=Exception('every time error'))
            Hoge().retry_with_error(mock3)
            self.assertEqual(mock3.execute.call_count, 3)


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

クラスをMockする場合

対象の実装

app/fuga.py

from domain

# domain.Slackオブジェクトを呼び出す実装
class Fuga:
    def notify(self) -> None:
        domain.Slack().notify('title', 'text')

テストの実装

from app import Fuga
import unittest
from unittest.mock import patch

class TestFuga(unittest.TestCase):

    def test_notify(self):
        # patch 定義で内部生成されるオブジェクトをmock化できる
        with patch('domain.Slack.notify', return_value=None) as mock:
            Fuga().notify()
            self.assertEqual(mock.call_count, 1)

参照