|  | 
| 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