fix(tests): fix: tests which call HTTP endpoints directly with the example parameters

Author: stainless-app[bot]

tests/test_client.py

@@ -23,9 +23,7 @@
 
 from replicate import Replicate, AsyncReplicate, APIResponseValidationError
 from replicate._types import Omit
-from replicate._utils import maybe_transform
 from replicate._models import BaseModel, FinalRequestOptions
-from replicate._constants import RAW_RESPONSE_HEADER
 from replicate._exceptions import APIStatusError, ReplicateError, APITimeoutError, APIResponseValidationError
 from replicate._base_client import (
     DEFAULT_TIMEOUT,
@@ -35,7 +33,6 @@
     DefaultAsyncHttpxClient,
     make_request_options,
 )
-from replicate.types.prediction_create_params import PredictionCreateParams
 
 from .utils import update_env
 
@@ -743,50 +740,27 @@ def test_parse_retry_after_header(self, remaining_retries: int, retry_after: str
 
     @mock.patch("replicate._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
     @pytest.mark.respx(base_url=base_url)
-    def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None:
+    def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: Replicate) -> None:
         respx_mock.post("/predictions").mock(side_effect=httpx.TimeoutException("Test timeout error"))
 
         with pytest.raises(APITimeoutError):
-            self.client.post(
-                "/predictions",
-                body=cast(
-                    object,
-                    maybe_transform(
-                        dict(
-                            input={"text": "Alice"},
-                            version="replicate/hello-world:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
-                        ),
-                        PredictionCreateParams,
-                    ),
-                ),
-                cast_to=httpx.Response,
-                options={"headers": {RAW_RESPONSE_HEADER: "stream"}},
-            )
+            client.predictions.with_streaming_response.create(
+                input={"text": "Alice"},
+                version="replicate/hello-world:9dcd6d78e7c6560c340d916fe32e9f24aabfa331e5cce95fe31f77fb03121426",
+            ).__enter__()
 
         assert _get_open_connections(self.client) == 0
 
     @mock.patch("replicate._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
     @pytest.mark.respx(base_url=base_url)
-    def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None:
+    def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: Replicate) -> None:
         respx_mock.post("/predictions").mock(return_value=httpx.Response(500))
 
         with pytest.raises(APIStatusError):
-            self.client.post(
-                "/predictions",
-                body=cast(
-                    object,
-                    maybe_transform(
-                        dict(
-                            input={"text": "Alice"},
-                            version="replicate/hello-world:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
-                        ),
-                        PredictionCreateParams,
-                    ),
-                ),
-                cast_to=httpx.Response,
-                options={"headers": {RAW_RESPONSE_HEADER: "stream"}},
-            )
-
+            client.predictions.with_streaming_response.create(
+                input={"text": "Alice"},
+                version="replicate/hello-world:9dcd6d78e7c6560c340d916fe32e9f24aabfa331e5cce95fe31f77fb03121426",
+            ).__enter__()
         assert _get_open_connections(self.client) == 0
 
     @pytest.mark.parametrize("failures_before_success", [0, 2, 4])
@@ -1615,50 +1589,31 @@ async def test_parse_retry_after_header(self, remaining_retries: int, retry_afte
 
     @mock.patch("replicate._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
     @pytest.mark.respx(base_url=base_url)
-    async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None:
+    async def test_retrying_timeout_errors_doesnt_leak(
+        self, respx_mock: MockRouter, async_client: AsyncReplicate
+    ) -> None:
         respx_mock.post("/predictions").mock(side_effect=httpx.TimeoutException("Test timeout error"))
 
         with pytest.raises(APITimeoutError):
-            await self.client.post(
-                "/predictions",
-                body=cast(
-                    object,
-                    maybe_transform(
-                        dict(
-                            input={"text": "Alice"},
-                            version="replicate/hello-world:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
-                        ),
-                        PredictionCreateParams,
-                    ),
-                ),
-                cast_to=httpx.Response,
-                options={"headers": {RAW_RESPONSE_HEADER: "stream"}},
-            )
+            await async_client.predictions.with_streaming_response.create(
+                input={"text": "Alice"},
+                version="replicate/hello-world:9dcd6d78e7c6560c340d916fe32e9f24aabfa331e5cce95fe31f77fb03121426",
+            ).__aenter__()
 
         assert _get_open_connections(self.client) == 0
 
     @mock.patch("replicate._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout)
     @pytest.mark.respx(base_url=base_url)
-    async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None:
+    async def test_retrying_status_errors_doesnt_leak(
+        self, respx_mock: MockRouter, async_client: AsyncReplicate
+    ) -> None:
         respx_mock.post("/predictions").mock(return_value=httpx.Response(500))
 
         with pytest.raises(APIStatusError):
-            await self.client.post(
-                "/predictions",
-                body=cast(
-                    object,
-                    maybe_transform(
-                        dict(
-                            input={"text": "Alice"},
-                            version="replicate/hello-world:5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa",
-                        ),
-                        PredictionCreateParams,
-                    ),
-                ),
-                cast_to=httpx.Response,
-                options={"headers": {RAW_RESPONSE_HEADER: "stream"}},
-            )
-
+            await async_client.predictions.with_streaming_response.create(
+                input={"text": "Alice"},
+                version="replicate/hello-world:9dcd6d78e7c6560c340d916fe32e9f24aabfa331e5cce95fe31f77fb03121426",
+            ).__aenter__()
         assert _get_open_connections(self.client) == 0
 
     @pytest.mark.parametrize("failures_before_success", [0, 2, 4])