|
23 | 23 |
|
24 | 24 | from llama_stack_client import LlamaStackClient, AsyncLlamaStackClient, APIResponseValidationError |
25 | 25 | from llama_stack_client._types import Omit |
26 | | -from llama_stack_client._utils import maybe_transform |
27 | 26 | from llama_stack_client._models import BaseModel, FinalRequestOptions |
28 | | -from llama_stack_client._constants import RAW_RESPONSE_HEADER |
29 | 27 | from llama_stack_client._exceptions import APIStatusError, APITimeoutError, APIResponseValidationError |
30 | 28 | from llama_stack_client._base_client import ( |
31 | 29 | DEFAULT_TIMEOUT, |
|
35 | 33 | DefaultAsyncHttpxClient, |
36 | 34 | make_request_options, |
37 | 35 | ) |
38 | | -from llama_stack_client.types.datasetio_append_rows_params import DatasetioAppendRowsParams |
39 | 36 |
|
40 | 37 | from .utils import update_env |
41 | 38 |
|
@@ -734,34 +731,27 @@ def test_parse_retry_after_header(self, remaining_retries: int, retry_after: str |
734 | 731 |
|
735 | 732 | @mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
736 | 733 | @pytest.mark.respx(base_url=base_url) |
737 | | - def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 734 | + def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: LlamaStackClient) -> None: |
738 | 735 | respx_mock.post("/v1/datasetio/append-rows/dataset_id").mock( |
739 | 736 | side_effect=httpx.TimeoutException("Test timeout error") |
740 | 737 | ) |
741 | 738 |
|
742 | 739 | with pytest.raises(APITimeoutError): |
743 | | - self.client.post( |
744 | | - "/v1/datasetio/append-rows/dataset_id", |
745 | | - body=cast(object, maybe_transform(dict(rows=[{"foo": True}]), DatasetioAppendRowsParams)), |
746 | | - cast_to=httpx.Response, |
747 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
748 | | - ) |
| 740 | + client.datasetio.with_streaming_response.append_rows( |
| 741 | + dataset_id="dataset_id", rows=[{"foo": True}] |
| 742 | + ).__enter__() |
749 | 743 |
|
750 | 744 | assert _get_open_connections(self.client) == 0 |
751 | 745 |
|
752 | 746 | @mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
753 | 747 | @pytest.mark.respx(base_url=base_url) |
754 | | - def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 748 | + def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: LlamaStackClient) -> None: |
755 | 749 | respx_mock.post("/v1/datasetio/append-rows/dataset_id").mock(return_value=httpx.Response(500)) |
756 | 750 |
|
757 | 751 | with pytest.raises(APIStatusError): |
758 | | - self.client.post( |
759 | | - "/v1/datasetio/append-rows/dataset_id", |
760 | | - body=cast(object, maybe_transform(dict(rows=[{"foo": True}]), DatasetioAppendRowsParams)), |
761 | | - cast_to=httpx.Response, |
762 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
763 | | - ) |
764 | | - |
| 752 | + client.datasetio.with_streaming_response.append_rows( |
| 753 | + dataset_id="dataset_id", rows=[{"foo": True}] |
| 754 | + ).__enter__() |
765 | 755 | assert _get_open_connections(self.client) == 0 |
766 | 756 |
|
767 | 757 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
@@ -1574,34 +1564,31 @@ async def test_parse_retry_after_header(self, remaining_retries: int, retry_afte |
1574 | 1564 |
|
1575 | 1565 | @mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1576 | 1566 | @pytest.mark.respx(base_url=base_url) |
1577 | | - async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1567 | + async def test_retrying_timeout_errors_doesnt_leak( |
| 1568 | + self, respx_mock: MockRouter, async_client: AsyncLlamaStackClient |
| 1569 | + ) -> None: |
1578 | 1570 | respx_mock.post("/v1/datasetio/append-rows/dataset_id").mock( |
1579 | 1571 | side_effect=httpx.TimeoutException("Test timeout error") |
1580 | 1572 | ) |
1581 | 1573 |
|
1582 | 1574 | with pytest.raises(APITimeoutError): |
1583 | | - await self.client.post( |
1584 | | - "/v1/datasetio/append-rows/dataset_id", |
1585 | | - body=cast(object, maybe_transform(dict(rows=[{"foo": True}]), DatasetioAppendRowsParams)), |
1586 | | - cast_to=httpx.Response, |
1587 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1588 | | - ) |
| 1575 | + await async_client.datasetio.with_streaming_response.append_rows( |
| 1576 | + dataset_id="dataset_id", rows=[{"foo": True}] |
| 1577 | + ).__aenter__() |
1589 | 1578 |
|
1590 | 1579 | assert _get_open_connections(self.client) == 0 |
1591 | 1580 |
|
1592 | 1581 | @mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1593 | 1582 | @pytest.mark.respx(base_url=base_url) |
1594 | | - async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1583 | + async def test_retrying_status_errors_doesnt_leak( |
| 1584 | + self, respx_mock: MockRouter, async_client: AsyncLlamaStackClient |
| 1585 | + ) -> None: |
1595 | 1586 | respx_mock.post("/v1/datasetio/append-rows/dataset_id").mock(return_value=httpx.Response(500)) |
1596 | 1587 |
|
1597 | 1588 | with pytest.raises(APIStatusError): |
1598 | | - await self.client.post( |
1599 | | - "/v1/datasetio/append-rows/dataset_id", |
1600 | | - body=cast(object, maybe_transform(dict(rows=[{"foo": True}]), DatasetioAppendRowsParams)), |
1601 | | - cast_to=httpx.Response, |
1602 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1603 | | - ) |
1604 | | - |
| 1589 | + await async_client.datasetio.with_streaming_response.append_rows( |
| 1590 | + dataset_id="dataset_id", rows=[{"foo": True}] |
| 1591 | + ).__aenter__() |
1605 | 1592 | assert _get_open_connections(self.client) == 0 |
1606 | 1593 |
|
1607 | 1594 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
|
0 commit comments