Skip to content
This repository was archived by the owner on Aug 14, 2025. It is now read-only.

Commit 347a4bf

Browse files
fix(tests): fix: tests which call HTTP endpoints directly with the example parameters
1 parent 9b63e1b commit 347a4bf

File tree

1 file changed

+20
-33
lines changed

1 file changed

+20
-33
lines changed

tests/test_client.py

Lines changed: 20 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,7 @@
2323

2424
from llama_stack_client import LlamaStackClient, AsyncLlamaStackClient, APIResponseValidationError
2525
from llama_stack_client._types import Omit
26-
from llama_stack_client._utils import maybe_transform
2726
from llama_stack_client._models import BaseModel, FinalRequestOptions
28-
from llama_stack_client._constants import RAW_RESPONSE_HEADER
2927
from llama_stack_client._exceptions import APIStatusError, APITimeoutError, APIResponseValidationError
3028
from llama_stack_client._base_client import (
3129
DEFAULT_TIMEOUT,
@@ -35,7 +33,6 @@
3533
DefaultAsyncHttpxClient,
3634
make_request_options,
3735
)
38-
from llama_stack_client.types.datasetio_append_rows_params import DatasetioAppendRowsParams
3936

4037
from .utils import update_env
4138

@@ -734,34 +731,27 @@ def test_parse_retry_after_header(self, remaining_retries: int, retry_after: str
734731

735732
@mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
736733
@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:
738735
respx_mock.post("/v1/datasetio/append-rows/dataset_id").mock(
739736
side_effect=httpx.TimeoutException("Test timeout error")
740737
)
741738

742739
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__()
749743

750744
assert _get_open_connections(self.client) == 0
751745

752746
@mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
753747
@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:
755749
respx_mock.post("/v1/datasetio/append-rows/dataset_id").mock(return_value=httpx.Response(500))
756750

757751
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__()
765755
assert _get_open_connections(self.client) == 0
766756

767757
@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
15741564

15751565
@mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
15761566
@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:
15781570
respx_mock.post("/v1/datasetio/append-rows/dataset_id").mock(
15791571
side_effect=httpx.TimeoutException("Test timeout error")
15801572
)
15811573

15821574
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__()
15891578

15901579
assert _get_open_connections(self.client) == 0
15911580

15921581
@mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
15931582
@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:
15951586
respx_mock.post("/v1/datasetio/append-rows/dataset_id").mock(return_value=httpx.Response(500))
15961587

15971588
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__()
16051592
assert _get_open_connections(self.client) == 0
16061593

16071594
@pytest.mark.parametrize("failures_before_success", [0, 2, 4])

0 commit comments

Comments
 (0)