Skip to content

Commit 729c66e

Browse files
feat(client): add follow_redirects request option
1 parent 5dc8d89 commit 729c66e

File tree

4 files changed

+64
-0
lines changed

4 files changed

+64
-0
lines changed

src/rizaio/_base_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -960,6 +960,9 @@ def request(
960960
if self.custom_auth is not None:
961961
kwargs["auth"] = self.custom_auth
962962

963+
if options.follow_redirects is not None:
964+
kwargs["follow_redirects"] = options.follow_redirects
965+
963966
log.debug("Sending HTTP Request: %s %s", request.method, request.url)
964967

965968
response = None
@@ -1460,6 +1463,9 @@ async def request(
14601463
if self.custom_auth is not None:
14611464
kwargs["auth"] = self.custom_auth
14621465

1466+
if options.follow_redirects is not None:
1467+
kwargs["follow_redirects"] = options.follow_redirects
1468+
14631469
log.debug("Sending HTTP Request: %s %s", request.method, request.url)
14641470

14651471
response = None

src/rizaio/_models.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -737,6 +737,7 @@ class FinalRequestOptionsInput(TypedDict, total=False):
737737
idempotency_key: str
738738
json_data: Body
739739
extra_json: AnyMapping
740+
follow_redirects: bool
740741

741742

742743
@final
@@ -750,6 +751,7 @@ class FinalRequestOptions(pydantic.BaseModel):
750751
files: Union[HttpxRequestFiles, None] = None
751752
idempotency_key: Union[str, None] = None
752753
post_parser: Union[Callable[[Any], Any], NotGiven] = NotGiven()
754+
follow_redirects: Union[bool, None] = None
753755

754756
# It should be noted that we cannot use `json` here as that would override
755757
# a BaseModel method in an incompatible fashion.

src/rizaio/_types.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ class RequestOptions(TypedDict, total=False):
100100
params: Query
101101
extra_json: AnyMapping
102102
idempotency_key: str
103+
follow_redirects: bool
103104

104105

105106
# Sentinel class used until PEP 0661 is accepted
@@ -215,3 +216,4 @@ class _GenericAlias(Protocol):
215216

216217
class HttpxSendArgs(TypedDict, total=False):
217218
auth: httpx.Auth
219+
follow_redirects: bool

tests/test_client.py

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -811,6 +811,33 @@ def retry_handler(_request: httpx.Request) -> httpx.Response:
811811

812812
assert response.http_request.headers.get("x-stainless-retry-count") == "42"
813813

814+
@pytest.mark.respx(base_url=base_url)
815+
def test_follow_redirects(self, respx_mock: MockRouter) -> None:
816+
# Test that the default follow_redirects=True allows following redirects
817+
respx_mock.post("/redirect").mock(
818+
return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
819+
)
820+
respx_mock.get("/redirected").mock(return_value=httpx.Response(200, json={"status": "ok"}))
821+
822+
response = self.client.post("/redirect", body={"key": "value"}, cast_to=httpx.Response)
823+
assert response.status_code == 200
824+
assert response.json() == {"status": "ok"}
825+
826+
@pytest.mark.respx(base_url=base_url)
827+
def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
828+
# Test that follow_redirects=False prevents following redirects
829+
respx_mock.post("/redirect").mock(
830+
return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
831+
)
832+
833+
with pytest.raises(APIStatusError) as exc_info:
834+
self.client.post(
835+
"/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response
836+
)
837+
838+
assert exc_info.value.response.status_code == 302
839+
assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
840+
814841

815842
class TestAsyncRiza:
816843
client = AsyncRiza(base_url=base_url, api_key=api_key, _strict_response_validation=True)
@@ -1632,3 +1659,30 @@ async def test_main() -> None:
16321659
raise AssertionError("calling get_platform using asyncify resulted in a hung process")
16331660

16341661
time.sleep(0.1)
1662+
1663+
@pytest.mark.respx(base_url=base_url)
1664+
async def test_follow_redirects(self, respx_mock: MockRouter) -> None:
1665+
# Test that the default follow_redirects=True allows following redirects
1666+
respx_mock.post("/redirect").mock(
1667+
return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
1668+
)
1669+
respx_mock.get("/redirected").mock(return_value=httpx.Response(200, json={"status": "ok"}))
1670+
1671+
response = await self.client.post("/redirect", body={"key": "value"}, cast_to=httpx.Response)
1672+
assert response.status_code == 200
1673+
assert response.json() == {"status": "ok"}
1674+
1675+
@pytest.mark.respx(base_url=base_url)
1676+
async def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
1677+
# Test that follow_redirects=False prevents following redirects
1678+
respx_mock.post("/redirect").mock(
1679+
return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
1680+
)
1681+
1682+
with pytest.raises(APIStatusError) as exc_info:
1683+
await self.client.post(
1684+
"/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response
1685+
)
1686+
1687+
assert exc_info.value.response.status_code == 302
1688+
assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"

0 commit comments

Comments
 (0)