Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion nonebot/drivers/aiohttp.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,11 +191,26 @@ async def stream_request(
timeout=timeout,
) as response:
response_headers = response.headers.copy()
# aiohttp does not guarantee fixed-size chunks; re-chunk to exact size
buffer = bytearray()
async for chunk in response.content.iter_chunked(chunk_size):
if not chunk:
continue
buffer.extend(chunk)
while len(buffer) >= chunk_size:
out = bytes(buffer[:chunk_size])
del buffer[:chunk_size]
yield Response(
response.status,
headers=response_headers,
content=out,
request=setup,
)
if buffer:
yield Response(
response.status,
headers=response_headers,
content=chunk,
content=bytes(buffer),
request=setup,
)

Expand Down
41 changes: 41 additions & 0 deletions tests/test_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
WebSocketClientMixin,
WebSocketServerSetup,
)
from nonebot.drivers.aiohttp import Session as AiohttpSession
from nonebot.drivers.aiohttp import WebSocket as AiohttpWebSocket
from nonebot.exception import WebSocketClosed
from nonebot.params import Depends
Expand Down Expand Up @@ -596,6 +597,46 @@ async def test_http_client_session(driver: Driver, server_url: URL):
await anyio.sleep(1)


@pytest.mark.anyio
async def test_aiohttp_stream_request_skip_empty_chunk() -> None:
class _FakeContent:
async def iter_chunked(self, _: int):
for chunk in (b"ab", b"", b"cd", b"e"):
yield chunk

class _FakeResponse:
def __init__(self) -> None:
self.status = 200
self.headers = {"x-test": "1"}
self.content = _FakeContent()

class _FakeRequestContext:
async def __aenter__(self) -> _FakeResponse:
return _FakeResponse()

async def __aexit__(self, *args: object) -> bool:
return False

class _FakeClient:
def request(self, *args: object, **kwargs: object) -> _FakeRequestContext:
return _FakeRequestContext()

session = AiohttpSession()
session._client = _FakeClient() # type: ignore[assignment]

chunks = []
async for resp in session.stream_request(
Request("GET", "https://example.com"), chunk_size=2
):
assert resp.status_code == 200
assert resp.content
chunks.append(resp.content)

assert chunks == [b"ab", b"cd", b"e"]
assert b"".join(chunks) == b"abcde"
assert all(len(chunk) == 2 for chunk in chunks[:-1])


@pytest.mark.anyio
@pytest.mark.parametrize(
"driver",
Expand Down
Loading