[问题] 如何正确的 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