|
24 | 24 | from gradientai import GradientAI, AsyncGradientAI, APIResponseValidationError |
25 | 25 | from gradientai._types import Omit |
26 | 26 | from gradientai._models import BaseModel, FinalRequestOptions |
27 | | -from gradientai._constants import RAW_RESPONSE_HEADER |
28 | 27 | from gradientai._exceptions import APIStatusError, APITimeoutError, GradientAIError, APIResponseValidationError |
29 | 28 | from gradientai._base_client import ( |
30 | 29 | DEFAULT_TIMEOUT, |
@@ -721,30 +720,21 @@ def test_parse_retry_after_header(self, remaining_retries: int, retry_after: str |
721 | 720 |
|
722 | 721 | @mock.patch("gradientai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
723 | 722 | @pytest.mark.respx(base_url=base_url) |
724 | | - def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 723 | + def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: GradientAI) -> None: |
725 | 724 | respx_mock.get("/v2/gen-ai/agents/uuid/versions").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
726 | 725 |
|
727 | 726 | with pytest.raises(APITimeoutError): |
728 | | - self.client.get( |
729 | | - "/v2/gen-ai/agents/uuid/versions", |
730 | | - cast_to=httpx.Response, |
731 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
732 | | - ) |
| 727 | + client.agents.versions.with_streaming_response.list(uuid="uuid").__enter__() |
733 | 728 |
|
734 | 729 | assert _get_open_connections(self.client) == 0 |
735 | 730 |
|
736 | 731 | @mock.patch("gradientai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
737 | 732 | @pytest.mark.respx(base_url=base_url) |
738 | | - def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 733 | + def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: GradientAI) -> None: |
739 | 734 | respx_mock.get("/v2/gen-ai/agents/uuid/versions").mock(return_value=httpx.Response(500)) |
740 | 735 |
|
741 | 736 | with pytest.raises(APIStatusError): |
742 | | - self.client.get( |
743 | | - "/v2/gen-ai/agents/uuid/versions", |
744 | | - cast_to=httpx.Response, |
745 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
746 | | - ) |
747 | | - |
| 737 | + client.agents.versions.with_streaming_response.list(uuid="uuid").__enter__() |
748 | 738 | assert _get_open_connections(self.client) == 0 |
749 | 739 |
|
750 | 740 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
@@ -1548,30 +1538,25 @@ async def test_parse_retry_after_header(self, remaining_retries: int, retry_afte |
1548 | 1538 |
|
1549 | 1539 | @mock.patch("gradientai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1550 | 1540 | @pytest.mark.respx(base_url=base_url) |
1551 | | - async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1541 | + async def test_retrying_timeout_errors_doesnt_leak( |
| 1542 | + self, respx_mock: MockRouter, async_client: AsyncGradientAI |
| 1543 | + ) -> None: |
1552 | 1544 | respx_mock.get("/v2/gen-ai/agents/uuid/versions").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
1553 | 1545 |
|
1554 | 1546 | with pytest.raises(APITimeoutError): |
1555 | | - await self.client.get( |
1556 | | - "/v2/gen-ai/agents/uuid/versions", |
1557 | | - cast_to=httpx.Response, |
1558 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1559 | | - ) |
| 1547 | + await async_client.agents.versions.with_streaming_response.list(uuid="uuid").__aenter__() |
1560 | 1548 |
|
1561 | 1549 | assert _get_open_connections(self.client) == 0 |
1562 | 1550 |
|
1563 | 1551 | @mock.patch("gradientai._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1564 | 1552 | @pytest.mark.respx(base_url=base_url) |
1565 | | - async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1553 | + async def test_retrying_status_errors_doesnt_leak( |
| 1554 | + self, respx_mock: MockRouter, async_client: AsyncGradientAI |
| 1555 | + ) -> None: |
1566 | 1556 | respx_mock.get("/v2/gen-ai/agents/uuid/versions").mock(return_value=httpx.Response(500)) |
1567 | 1557 |
|
1568 | 1558 | with pytest.raises(APIStatusError): |
1569 | | - await self.client.get( |
1570 | | - "/v2/gen-ai/agents/uuid/versions", |
1571 | | - cast_to=httpx.Response, |
1572 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1573 | | - ) |
1574 | | - |
| 1559 | + await async_client.agents.versions.with_streaming_response.list(uuid="uuid").__aenter__() |
1575 | 1560 | assert _get_open_connections(self.client) == 0 |
1576 | 1561 |
|
1577 | 1562 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
|
0 commit comments