diff --git a/nonebot/drivers/aiohttp.py b/nonebot/drivers/aiohttp.py index bc063cda0bb1..65a30afb5307 100644 --- a/nonebot/drivers/aiohttp.py +++ b/nonebot/drivers/aiohttp.py @@ -169,6 +169,7 @@ async def request(self, setup: Request) -> Response: headers=setup.headers, proxy=setup.proxy or self._proxy, timeout=self._get_timeout(setup.timeout), + **exclude_unset({"allow_redirects": setup.auto_redirects}), ) as response: return Response( response.status, @@ -210,6 +211,7 @@ async def stream_request( headers=setup.headers, proxy=setup.proxy or self._proxy, timeout=self._get_timeout(setup.timeout), + **exclude_unset({"allow_redirects": setup.auto_redirects}), ) as response: response_headers = response.headers.copy() # aiohttp does not guarantee fixed-size chunks; re-chunk to exact size @@ -335,6 +337,12 @@ async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]: "the configured ping timeout will be ignored." ) + if setup.auto_redirects is not UNSET: + logger.warning( + "aiohttp driver websocket method does not expose redirects parameter; " + "the configured auto_redirects will be ignored." + ) + async with aiohttp.ClientSession(version=version, trust_env=True) as session: async with session.ws_connect( setup.url, diff --git a/nonebot/drivers/httpx.py b/nonebot/drivers/httpx.py index a8217b793c53..128b1d37500d 100644 --- a/nonebot/drivers/httpx.py +++ b/nonebot/drivers/httpx.py @@ -145,6 +145,7 @@ async def request(self, setup: Request) -> Response: headers=tuple(setup.headers.items()), cookies=setup.cookies.jar, timeout=self._get_timeout(setup.timeout), + **exclude_unset({"follow_redirects": setup.auto_redirects}), ) return Response( response.status_code, @@ -172,6 +173,7 @@ async def stream_request( headers=tuple(setup.headers.items()), cookies=setup.cookies.jar, timeout=self._get_timeout(setup.timeout), + **exclude_unset({"follow_redirects": setup.auto_redirects}), ) as response: response_headers = response.headers.multi_items() async for chunk in response.aiter_bytes(chunk_size=chunk_size): diff --git a/nonebot/drivers/httpx2.py b/nonebot/drivers/httpx2.py index a190d49bd59d..b78075630b4e 100644 --- a/nonebot/drivers/httpx2.py +++ b/nonebot/drivers/httpx2.py @@ -145,6 +145,7 @@ async def request(self, setup: Request) -> Response: headers=tuple(setup.headers.items()), cookies=setup.cookies.jar, timeout=self._get_timeout(setup.timeout), + **exclude_unset({"follow_redirects": setup.auto_redirects}), ) return Response( response.status_code, @@ -172,6 +173,7 @@ async def stream_request( headers=tuple(setup.headers.items()), cookies=setup.cookies.jar, timeout=self._get_timeout(setup.timeout), + **exclude_unset({"follow_redirects": setup.auto_redirects}), ) as response: response_headers = response.headers.multi_items() async for chunk in response.aiter_bytes(chunk_size=chunk_size): diff --git a/nonebot/drivers/websockets.py b/nonebot/drivers/websockets.py index 0580d6eee24c..d1dde5fcf39f 100644 --- a/nonebot/drivers/websockets.py +++ b/nonebot/drivers/websockets.py @@ -111,6 +111,12 @@ async def websocket(self, setup: Request) -> AsyncGenerator["WebSocket", None]: } ) + if setup.auto_redirects is not UNSET: + logger.warning( + "websockets driver websocket method does not expose " + "redirects parameter; the configured auto_redirects will be ignored." + ) + connection = connect( str(setup.url), additional_headers={**setup.headers, **setup.cookies.as_header(setup)}, diff --git a/nonebot/internal/driver/model.py b/nonebot/internal/driver/model.py index f4fb0b4a9d8f..7385557ff6d2 100644 --- a/nonebot/internal/driver/model.py +++ b/nonebot/internal/driver/model.py @@ -79,6 +79,7 @@ def __init__( timeout: TimeoutTypes | UnsetType = UNSET, proxy: str | None = None, ping_interval: PingIntervalTypes | UnsetType = UNSET, + auto_redirects: bool | UnsetType = UNSET, ): # method self.method: str = ( @@ -94,6 +95,8 @@ def __init__( self.proxy: str | None = proxy # ping interval self.ping_interval: PingIntervalTypes | UnsetType = ping_interval + # auto redirects + self.auto_redirects: bool | UnsetType = auto_redirects # url if isinstance(url, tuple): diff --git a/tests/fake_server.py b/tests/fake_server.py index ea473286675c..bafacb4b80ad 100644 --- a/tests/fake_server.py +++ b/tests/fake_server.py @@ -79,6 +79,10 @@ def http_echo(request: Request) -> Response: ) +def http_redirect(request: Request) -> Response: + return Response("redirecting", status=302, headers={"Location": "/"}) + + def websocket_echo(request: Request) -> Response: stream = request.environ["werkzeug.socket"] @@ -137,5 +141,7 @@ def websocket_echo(request: Request) -> Response: def request_handler(request: Request) -> Response: if request.headers.get("Connection") == "Upgrade": return websocket_echo(request) + elif request.path == "/redirect": + return http_redirect(request) else: return http_echo(request) diff --git a/tests/test_driver.py b/tests/test_driver.py index 57930a4cb463..71973ed31c3e 100644 --- a/tests/test_driver.py +++ b/tests/test_driver.py @@ -372,8 +372,8 @@ async def test_http_client(driver: Driver, server_url: URL): ) chunks = [] async for resp in driver.stream_request(request, chunk_size=4): - assert response.status_code == 200 - assert response.content + assert resp.status_code == 200 + assert resp.content chunks.append(resp.content) assert all(len(chunk) == 4 for chunk in chunks[:-1]) data = json.loads(b"".join(chunks)) @@ -388,6 +388,62 @@ async def test_http_client(driver: Driver, server_url: URL): await anyio.sleep(1) +@pytest.mark.anyio +@pytest.mark.parametrize( + "driver", + [ + pytest.param("nonebot.drivers.httpx:Driver", id="httpx"), + pytest.param("nonebot.drivers.httpx2:Driver", id="httpx2"), + pytest.param("nonebot.drivers.aiohttp:Driver", id="aiohttp"), + ], + indirect=True, +) +async def test_http_client_redirects(driver: Driver, server_url: URL): + assert isinstance(driver, HTTPClientMixin) + + redirect_url = server_url.join(URL("/redirect")) + + # follow redirects by default + request = Request("GET", redirect_url) + assert request.auto_redirects is UNSET + response = await driver.request(request) + assert response.status_code == 200 + assert response.content + data = json.loads(response.content) + assert data["method"] == "GET" + assert URL(data["url"]).path == "/" + + # disable auto redirects + request = Request("GET", redirect_url, auto_redirects=False) + assert request.auto_redirects is False + response = await driver.request(request) + assert response.status_code == 302 + assert response.headers["Location"] == "/" + assert response.content == b"redirecting" + + # stream request following redirects + request = Request("GET", redirect_url) + chunks = [] + async for resp in driver.stream_request(request, chunk_size=4): + assert resp.status_code == 200 + assert resp.content + chunks.append(resp.content) + data = json.loads(b"".join(chunks)) + assert data["method"] == "GET" + assert URL(data["url"]).path == "/" + + # stream request without auto redirects + request = Request("GET", redirect_url, auto_redirects=False) + chunks = [] + async for resp in driver.stream_request(request, chunk_size=4): + assert resp.status_code == 302 + assert resp.content + chunks.append(resp.content) + assert b"".join(chunks) == b"redirecting" + + await anyio.sleep(1) + + @pytest.mark.anyio @pytest.mark.parametrize( "driver", @@ -600,6 +656,64 @@ async def test_http_client_session(driver: Driver, server_url: URL): await anyio.sleep(1) +@pytest.mark.anyio +@pytest.mark.parametrize( + "driver", + [ + pytest.param("nonebot.drivers.httpx:Driver", id="httpx"), + pytest.param("nonebot.drivers.httpx2:Driver", id="httpx2"), + pytest.param("nonebot.drivers.aiohttp:Driver", id="aiohttp"), + ], + indirect=True, +) +async def test_http_client_session_redirects(driver: Driver, server_url: URL): + assert isinstance(driver, HTTPClientMixin) + + redirect_url = server_url.join(URL("/redirect")) + + session = driver.get_session() + async with session: + # follow redirects by default + request = Request("GET", redirect_url) + assert request.auto_redirects is UNSET + response = await session.request(request) + assert response.status_code == 200 + assert response.content + data = json.loads(response.content) + assert data["method"] == "GET" + assert URL(data["url"]).path == "/" + + # disable auto redirects + request = Request("GET", redirect_url, auto_redirects=False) + assert request.auto_redirects is False + response = await session.request(request) + assert response.status_code == 302 + assert response.headers["Location"] == "/" + assert response.content == b"redirecting" + + # stream request following redirects + request = Request("GET", redirect_url) + chunks = [] + async for resp in session.stream_request(request, chunk_size=4): + assert resp.status_code == 200 + assert resp.content + chunks.append(resp.content) + data = json.loads(b"".join(chunks)) + assert data["method"] == "GET" + assert URL(data["url"]).path == "/" + + # stream request without auto redirects + request = Request("GET", redirect_url, auto_redirects=False) + chunks = [] + async for resp in session.stream_request(request, chunk_size=4): + assert resp.status_code == 302 + assert resp.content + chunks.append(resp.content) + assert b"".join(chunks) == b"redirecting" + + await anyio.sleep(1) + + @pytest.mark.anyio async def test_aiohttp_stream_request_skip_empty_chunk() -> None: class _FakeContent: