|
23 | 23 |
|
24 | 24 | from asktable import Asktable, AsyncAsktable, APIResponseValidationError |
25 | 25 | from asktable._types import Omit |
26 | | -from asktable._utils import maybe_transform |
27 | 26 | from asktable._models import BaseModel, FinalRequestOptions |
28 | | -from asktable._constants import RAW_RESPONSE_HEADER |
29 | 27 | from asktable._exceptions import APIStatusError, APITimeoutError, APIResponseValidationError |
30 | 28 | from asktable._base_client import ( |
31 | 29 | DEFAULT_TIMEOUT, |
|
35 | 33 | DefaultAsyncHttpxClient, |
36 | 34 | make_request_options, |
37 | 35 | ) |
38 | | -from asktable.types.datasource_create_params import DatasourceCreateParams |
39 | 36 |
|
40 | 37 | from .utils import update_env |
41 | 38 |
|
@@ -705,32 +702,21 @@ def test_parse_retry_after_header(self, remaining_retries: int, retry_after: str |
705 | 702 |
|
706 | 703 | @mock.patch("asktable._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
707 | 704 | @pytest.mark.respx(base_url=base_url) |
708 | | - def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 705 | + def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: Asktable) -> None: |
709 | 706 | respx_mock.post("/v1/datasources").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
710 | 707 |
|
711 | 708 | with pytest.raises(APITimeoutError): |
712 | | - self.client.post( |
713 | | - "/v1/datasources", |
714 | | - body=cast(object, maybe_transform(dict(engine="mysql"), DatasourceCreateParams)), |
715 | | - cast_to=httpx.Response, |
716 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
717 | | - ) |
| 709 | + client.datasources.with_streaming_response.create(engine="mysql").__enter__() |
718 | 710 |
|
719 | 711 | assert _get_open_connections(self.client) == 0 |
720 | 712 |
|
721 | 713 | @mock.patch("asktable._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
722 | 714 | @pytest.mark.respx(base_url=base_url) |
723 | | - def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 715 | + def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: Asktable) -> None: |
724 | 716 | respx_mock.post("/v1/datasources").mock(return_value=httpx.Response(500)) |
725 | 717 |
|
726 | 718 | with pytest.raises(APIStatusError): |
727 | | - self.client.post( |
728 | | - "/v1/datasources", |
729 | | - body=cast(object, maybe_transform(dict(engine="mysql"), DatasourceCreateParams)), |
730 | | - cast_to=httpx.Response, |
731 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
732 | | - ) |
733 | | - |
| 719 | + client.datasources.with_streaming_response.create(engine="mysql").__enter__() |
734 | 720 | assert _get_open_connections(self.client) == 0 |
735 | 721 |
|
736 | 722 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
@@ -1524,32 +1510,25 @@ async def test_parse_retry_after_header(self, remaining_retries: int, retry_afte |
1524 | 1510 |
|
1525 | 1511 | @mock.patch("asktable._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1526 | 1512 | @pytest.mark.respx(base_url=base_url) |
1527 | | - async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1513 | + async def test_retrying_timeout_errors_doesnt_leak( |
| 1514 | + self, respx_mock: MockRouter, async_client: AsyncAsktable |
| 1515 | + ) -> None: |
1528 | 1516 | respx_mock.post("/v1/datasources").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
1529 | 1517 |
|
1530 | 1518 | with pytest.raises(APITimeoutError): |
1531 | | - await self.client.post( |
1532 | | - "/v1/datasources", |
1533 | | - body=cast(object, maybe_transform(dict(engine="mysql"), DatasourceCreateParams)), |
1534 | | - cast_to=httpx.Response, |
1535 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1536 | | - ) |
| 1519 | + await async_client.datasources.with_streaming_response.create(engine="mysql").__aenter__() |
1537 | 1520 |
|
1538 | 1521 | assert _get_open_connections(self.client) == 0 |
1539 | 1522 |
|
1540 | 1523 | @mock.patch("asktable._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1541 | 1524 | @pytest.mark.respx(base_url=base_url) |
1542 | | - async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1525 | + async def test_retrying_status_errors_doesnt_leak( |
| 1526 | + self, respx_mock: MockRouter, async_client: AsyncAsktable |
| 1527 | + ) -> None: |
1543 | 1528 | respx_mock.post("/v1/datasources").mock(return_value=httpx.Response(500)) |
1544 | 1529 |
|
1545 | 1530 | with pytest.raises(APIStatusError): |
1546 | | - await self.client.post( |
1547 | | - "/v1/datasources", |
1548 | | - body=cast(object, maybe_transform(dict(engine="mysql"), DatasourceCreateParams)), |
1549 | | - cast_to=httpx.Response, |
1550 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1551 | | - ) |
1552 | | - |
| 1531 | + await async_client.datasources.with_streaming_response.create(engine="mysql").__aenter__() |
1553 | 1532 | assert _get_open_connections(self.client) == 0 |
1554 | 1533 |
|
1555 | 1534 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
|
0 commit comments