|
23 | 23 |
|
24 | 24 | from zeroentropy import ZeroEntropy, AsyncZeroEntropy, APIResponseValidationError |
25 | 25 | from zeroentropy._types import Omit |
26 | | -from zeroentropy._utils import maybe_transform |
27 | 26 | from zeroentropy._models import BaseModel, FinalRequestOptions |
28 | | -from zeroentropy._constants import RAW_RESPONSE_HEADER |
29 | 27 | from zeroentropy._exceptions import APIStatusError, APITimeoutError, ZeroEntropyError, APIResponseValidationError |
30 | 28 | from zeroentropy._base_client import ( |
31 | 29 | DEFAULT_TIMEOUT, |
|
35 | 33 | DefaultAsyncHttpxClient, |
36 | 34 | make_request_options, |
37 | 35 | ) |
38 | | -from zeroentropy.types.status_get_status_params import StatusGetStatusParams |
39 | 36 |
|
40 | 37 | from .utils import update_env |
41 | 38 |
|
@@ -725,32 +722,21 @@ def test_parse_retry_after_header(self, remaining_retries: int, retry_after: str |
725 | 722 |
|
726 | 723 | @mock.patch("zeroentropy._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
727 | 724 | @pytest.mark.respx(base_url=base_url) |
728 | | - def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 725 | + def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: ZeroEntropy) -> None: |
729 | 726 | respx_mock.post("/status/get-status").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
730 | 727 |
|
731 | 728 | with pytest.raises(APITimeoutError): |
732 | | - self.client.post( |
733 | | - "/status/get-status", |
734 | | - body=cast(object, maybe_transform({}, StatusGetStatusParams)), |
735 | | - cast_to=httpx.Response, |
736 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
737 | | - ) |
| 729 | + client.status.with_streaming_response.get_status().__enter__() |
738 | 730 |
|
739 | 731 | assert _get_open_connections(self.client) == 0 |
740 | 732 |
|
741 | 733 | @mock.patch("zeroentropy._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
742 | 734 | @pytest.mark.respx(base_url=base_url) |
743 | | - def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 735 | + def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: ZeroEntropy) -> None: |
744 | 736 | respx_mock.post("/status/get-status").mock(return_value=httpx.Response(500)) |
745 | 737 |
|
746 | 738 | with pytest.raises(APIStatusError): |
747 | | - self.client.post( |
748 | | - "/status/get-status", |
749 | | - body=cast(object, maybe_transform({}, StatusGetStatusParams)), |
750 | | - cast_to=httpx.Response, |
751 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
752 | | - ) |
753 | | - |
| 739 | + client.status.with_streaming_response.get_status().__enter__() |
754 | 740 | assert _get_open_connections(self.client) == 0 |
755 | 741 |
|
756 | 742 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
@@ -1550,32 +1536,25 @@ async def test_parse_retry_after_header(self, remaining_retries: int, retry_afte |
1550 | 1536 |
|
1551 | 1537 | @mock.patch("zeroentropy._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1552 | 1538 | @pytest.mark.respx(base_url=base_url) |
1553 | | - async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1539 | + async def test_retrying_timeout_errors_doesnt_leak( |
| 1540 | + self, respx_mock: MockRouter, async_client: AsyncZeroEntropy |
| 1541 | + ) -> None: |
1554 | 1542 | respx_mock.post("/status/get-status").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
1555 | 1543 |
|
1556 | 1544 | with pytest.raises(APITimeoutError): |
1557 | | - await self.client.post( |
1558 | | - "/status/get-status", |
1559 | | - body=cast(object, maybe_transform({}, StatusGetStatusParams)), |
1560 | | - cast_to=httpx.Response, |
1561 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1562 | | - ) |
| 1545 | + await async_client.status.with_streaming_response.get_status().__aenter__() |
1563 | 1546 |
|
1564 | 1547 | assert _get_open_connections(self.client) == 0 |
1565 | 1548 |
|
1566 | 1549 | @mock.patch("zeroentropy._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1567 | 1550 | @pytest.mark.respx(base_url=base_url) |
1568 | | - async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1551 | + async def test_retrying_status_errors_doesnt_leak( |
| 1552 | + self, respx_mock: MockRouter, async_client: AsyncZeroEntropy |
| 1553 | + ) -> None: |
1569 | 1554 | respx_mock.post("/status/get-status").mock(return_value=httpx.Response(500)) |
1570 | 1555 |
|
1571 | 1556 | with pytest.raises(APIStatusError): |
1572 | | - await self.client.post( |
1573 | | - "/status/get-status", |
1574 | | - body=cast(object, maybe_transform({}, StatusGetStatusParams)), |
1575 | | - cast_to=httpx.Response, |
1576 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1577 | | - ) |
1578 | | - |
| 1557 | + await async_client.status.with_streaming_response.get_status().__aenter__() |
1579 | 1558 | assert _get_open_connections(self.client) == 0 |
1580 | 1559 |
|
1581 | 1560 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
|
0 commit comments