Python 模拟response.content.iter_chunked

Question

这是我的python 3.8代码，为什么我不能模拟iter_chucked？

async def _http_get(part_path):
    async with aiofiles.open(part_path, "ab+") as temp_file:
        async with RetryClient(client_session=ClientSession(),
                               retry_options=RandomRetry(statuses=[408, 429, 500, 502, 503, 504])) as session:
            timeout = aiohttp.ClientTimeout(total=3600)
            async with session.request(method="GET", url="http://mock.com",
                                       ssl=ssl.create_default_context(cafile=where()),
                                       timeout=timeout) as response:

                async for chunk in response.content.iter_chunked(10 * 1024 * 1024): # here what I want to mock
                    if chunk:  # filter out keep-alive new chunks
                        await temp_file.write(chunk)

这是我的测试代码

async def mock_iter_content_async(n):
    for chunk in [b'chunk1', b'chunk2', b'chunk3']:
        yield chunk

class TestFileDownload(unittest.TestCase):
    mock_response = MagicMock()
    mock_response.content = mock.Mock()
    mock_response.content.iter_chunked.return_value = mock_iter_content_async(11)

    @patch("aiohttp_retry.RetryClient.request", return_value=mock_response)
    def test_async_http(self, mockClass1):
        part_path = os.path.join(os.getcwd(), "parts.txt")
        loop = asyncio.get_event_loop()
        result = loop.run_until_complete(_http_get(part_path))

但出现错误

TypeError: 'async for' requires an object with __aiter__ method, got coroutine

我发现我模拟的不起作用，如何模拟iter_chunked？

Answer 1

你的模拟不起作用的原因是因为

iter_chunked

是一个协程方法，而

mock_iter_content_async

是一个异步生成器函数。它们不直接兼容使用 MagicMock 或 mock.Mock 进行模拟。

要正确模拟

iter_chunked

，您需要创建一个行为类似于异步迭代器的模拟。您可以通过使用

AsyncMock

模块提供的

unittest.mock

来实现此目的。以下是如何修改测试代码以正确模拟

iter_chunked

:

import asyncio
import aiofiles
import aiohttp
import ssl
import os
from aiohttp import ClientSession
from aiohttp_retry import RetryClient
from aiohttp_retry.retriable import RandomRetry
from unittest.mock import patch, AsyncMock
import unittest


async def _http_get(part_path):
    async with aiofiles.open(part_path, "ab+") as temp_file:
        async with RetryClient(client_session=ClientSession(),
                               retry_options=RandomRetry(statuses=[408, 429, 500, 502, 503, 504])) as session:
            timeout = aiohttp.ClientTimeout(total=3600)
            async with session.request(method="GET", url="http://mock.com",
                                       ssl=ssl.create_default_context(cafile=where()),
                                       timeout=timeout) as response:
                async for chunk in response.content.iter_chunked(10 * 1024 * 1024):  # here what I want to mock
                    if chunk:  # filter out keep-alive new chunks
                        await temp_file.write(chunk)


async def mock_iter_content_async(n):
    for chunk in [b'chunk1', b'chunk2', b'chunk3']:
        yield chunk


class TestFileDownload(unittest.TestCase):
    @patch("aiohttp.ClientResponse.content", new_callable=AsyncMock)
    async def test_async_http(self, mock_content):
        mock_content.iter_chunked.return_value = mock_iter_content_async(11)

        part_path = os.path.join(os.getcwd(), "parts.txt")
        result = await _http_get(part_path)
        # Assert your result if needed

在此修改版本中，

AsyncMock

用于为

content

的

ClientResponse

属性创建模拟，并且

iter_chunked

被正确模拟以返回由

mock_iter_content_async

创建的异步生成器。这应该可以解决您遇到的错误。