[問題] 如何正確的 mock requests 回傳的 iter_content

[chan15 (ChaN)]

def stream_download(url: str, body: dict, dest: str) -> None:
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    with requests.post(url, get_default_body() or body, stream=True) as r:
        r.raise_for_status()
        with open(dest, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)


這是 streaming 抓取檔按的 method。


    @patch('my_http.get_default_body')
    @patch('requests.post')
    def test_stream_download(self, mock_post, mock_get_default_body):
        url = 'https://example.com'
        body = {'key': 'value'}

        mock_get_default_body.return_value = body
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.iter_content.return_value = [b'chunk1', b'chunk2']
        mock_post.return_value = mock_response

        with tempfile.TemporaryDirectory() as dest_path:
            dest = os.path.join(dest_path, 'temp')
            my_http.stream_download(url, body, dest)
            mock_post.assert_called_once_with(url, body, stream=True)

            with open(dest, 'rb') as f:
                self.assertEqual(f.read(), b'chunk1chunk2')


這是我目前的測試 code，想法是用 tempfile 產生一個真正的內容去比對結果
但試過很多方法不管怎麼塞 iter_content 的 data
進到 stream_download 的 method 時都會抓不到
自然也不會跑該迴圈的 open write method，請問該怎麼寫才是正確的

--
※ 發信站: 批踢踢實業坊(ptt.cc), 來自: 59.120.18.61 (臺灣)
※ 文章網址: https://www.ptt.cc/bbs/Python/M.1679370847.A.831.html
※ 編輯: chan15 (59.120.18.61 臺灣), 03/21/2023 11:54:50
→ lycantrope: 還需要mock response的 __enter__
→ chan15: 成功了，謝謝喔