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

作者: chan15 (ChaN)   2023-03-21 11:54:05
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,請問該怎麼寫才是正確的
作者: lycantrope (阿寬)   2023-03-21 12:29:00
還需要mock response的 __enter__
作者: chan15 (ChaN)   2023-03-21 14:56:00
成功了,謝謝喔

Links booklink

Contact Us: admin [ a t ] ucptt.com