mirror of
https://github.com/nonebot/nonebot2.git
synced 2026-04-17 06:07:55 +00:00
🐛 Fix: 修复 aiohttp 流式响应分块大小不固定问题 (#3919)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
@@ -191,11 +191,26 @@ class Session(HTTPClientSession):
|
|||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
) as response:
|
) as response:
|
||||||
response_headers = response.headers.copy()
|
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):
|
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(
|
yield Response(
|
||||||
response.status,
|
response.status,
|
||||||
headers=response_headers,
|
headers=response_headers,
|
||||||
content=chunk,
|
content=bytes(buffer),
|
||||||
request=setup,
|
request=setup,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ from nonebot.drivers import (
|
|||||||
WebSocketClientMixin,
|
WebSocketClientMixin,
|
||||||
WebSocketServerSetup,
|
WebSocketServerSetup,
|
||||||
)
|
)
|
||||||
|
from nonebot.drivers.aiohttp import Session as AiohttpSession
|
||||||
from nonebot.drivers.aiohttp import WebSocket as AiohttpWebSocket
|
from nonebot.drivers.aiohttp import WebSocket as AiohttpWebSocket
|
||||||
from nonebot.exception import WebSocketClosed
|
from nonebot.exception import WebSocketClosed
|
||||||
from nonebot.params import Depends
|
from nonebot.params import Depends
|
||||||
@@ -596,6 +597,46 @@ async def test_http_client_session(driver: Driver, server_url: URL):
|
|||||||
await anyio.sleep(1)
|
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.anyio
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"driver",
|
"driver",
|
||||||
|
|||||||
Reference in New Issue
Block a user