Skip to content

Commit 5185843

Browse files
feat(client): add follow_redirects request option
1 parent c0cdb68 commit 5185843

File tree

4 files changed

+64
-0
lines changed

4 files changed

+64
-0
lines changed

src/orb/_base_client.py

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

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

966969
response = None
@@ -1475,6 +1478,9 @@ async def request(
14751478
if self.custom_auth is not None:
14761479
kwargs["auth"] = self.custom_auth
14771480

1481+
if options.follow_redirects is not None:
1482+
kwargs["follow_redirects"] = options.follow_redirects
1483+
14781484
log.debug("Sending HTTP Request: %s %s", request.method, request.url)
14791485

14801486
response = None

src/orb/_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/orb/_types.py

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

105106

106107
# Sentinel class used until PEP 0661 is accepted
@@ -217,3 +218,4 @@ class _GenericAlias(Protocol):
217218

218219
class HttpxSendArgs(TypedDict, total=False):
219220
auth: httpx.Auth
221+
follow_redirects: bool

tests/test_client.py

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -869,6 +869,33 @@ def retry_handler(_request: httpx.Request) -> httpx.Response:
869869
assert response.retries_taken == failures_before_success
870870
assert int(response.http_request.headers.get("x-stainless-retry-count")) == failures_before_success
871871

872+
@pytest.mark.respx(base_url=base_url)
873+
def test_follow_redirects(self, respx_mock: MockRouter) -> None:
874+
# Test that the default follow_redirects=True allows following redirects
875+
respx_mock.post("/redirect").mock(
876+
return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
877+
)
878+
respx_mock.get("/redirected").mock(return_value=httpx.Response(200, json={"status": "ok"}))
879+
880+
response = self.client.post("/redirect", body={"key": "value"}, cast_to=httpx.Response)
881+
assert response.status_code == 200
882+
assert response.json() == {"status": "ok"}
883+
884+
@pytest.mark.respx(base_url=base_url)
885+
def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
886+
# Test that follow_redirects=False prevents following redirects
887+
respx_mock.post("/redirect").mock(
888+
return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
889+
)
890+
891+
with pytest.raises(APIStatusError) as exc_info:
892+
self.client.post(
893+
"/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response
894+
)
895+
896+
assert exc_info.value.response.status_code == 302
897+
assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"
898+
872899

873900
class TestAsyncOrb:
874901
client = AsyncOrb(base_url=base_url, api_key=api_key, _strict_response_validation=True)
@@ -1751,3 +1778,30 @@ async def test_main() -> None:
17511778
raise AssertionError("calling get_platform using asyncify resulted in a hung process")
17521779

17531780
time.sleep(0.1)
1781+
1782+
@pytest.mark.respx(base_url=base_url)
1783+
async def test_follow_redirects(self, respx_mock: MockRouter) -> None:
1784+
# Test that the default follow_redirects=True allows following redirects
1785+
respx_mock.post("/redirect").mock(
1786+
return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
1787+
)
1788+
respx_mock.get("/redirected").mock(return_value=httpx.Response(200, json={"status": "ok"}))
1789+
1790+
response = await self.client.post("/redirect", body={"key": "value"}, cast_to=httpx.Response)
1791+
assert response.status_code == 200
1792+
assert response.json() == {"status": "ok"}
1793+
1794+
@pytest.mark.respx(base_url=base_url)
1795+
async def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
1796+
# Test that follow_redirects=False prevents following redirects
1797+
respx_mock.post("/redirect").mock(
1798+
return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
1799+
)
1800+
1801+
with pytest.raises(APIStatusError) as exc_info:
1802+
await self.client.post(
1803+
"/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response
1804+
)
1805+
1806+
assert exc_info.value.response.status_code == 302
1807+
assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"

0 commit comments

Comments
 (0)