pythonのmockでwith句を使ったテストを実行する

mockを使ったテストを行っていて、検索してもなかなか解決策が見つからなかったので投稿。

結局公式ページから推測して結論を出しました。

With句で使うClass

まずはテスト対象となるWith句を使える形にしたclassのサンプル

sampleblock.py

class SampleBlock():
    def __enter__(self):
        pass

    def get_hoge(self):
        return "hoge"

    def __exit__(self, exc_type, exc_value, tb):
        pass

そして使う側のサンプル。
with句でくくりながらもそれ自身にメソッドが存在して、引数によってそれが実行されるかどうかが決まるみたいな使われ方。

runner.py

import sampleblock as sb

def hogehoge(flg):
    with sb.SampleBlock() as sampleblock:
        if flg:
            sampleblock.get_hoge()

実際のテストコード

引数によって、このsampleblockのget_hogeが呼び出されたかどうかを確認するチェックを行ったときどうするか？という問題でした。
もう一つ条件があって、enterとかexitに通信関係の仕組みが入っていてそこはmockにしておきたいというわがままの叶えたい。

test_hogehoge.py

import unittest
from unittest.mock import patch 

import runner

class HogeHogeTest(unittest.TestCase):
    def test_true(self):
        with patch('sampleblock.SampleBlock') as mock_obj:
            runner.hogehoge(true)
            self.assertTrue(mock_obj.return_value.__enter__.return_valuecalled)

    def test_false(self):
        with patch('sampleblock.SampleBlock') as mock_obj:
            runner.hogehoge(false)
            self.assertFalse(mock_obj.return_value.__enter__.return_valuecalled)

callをいったん出力して導き出した答え。
ちょっとこれはわかりづらい。with句特有の動きなのかどうかは検証が必要ですね。
実際のプログラムを書き換えて投稿しているので、間違っていたら指摘お願いします。