|
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, |
32 | 30 | HTTPX_DEFAULT_TIMEOUT, |
33 | 31 | BaseClient, |
| 32 | + DefaultHttpxClient, |
| 33 | + DefaultAsyncHttpxClient, |
34 | 34 | make_request_options, |
35 | 35 | ) |
36 | | -from llama_stack_client.types.inference_chat_completion_params import InferenceChatCompletionParamsNonStreaming |
37 | 36 |
|
38 | 37 | from .utils import update_env |
39 | 38 |
|
@@ -677,60 +676,37 @@ def test_parse_retry_after_header(self, remaining_retries: int, retry_after: str |
677 | 676 |
|
678 | 677 | @mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
679 | 678 | @pytest.mark.respx(base_url=base_url) |
680 | | - def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 679 | + def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: LlamaStackClient) -> None: |
681 | 680 | respx_mock.post("/v1/inference/chat-completion").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
682 | 681 |
|
683 | 682 | with pytest.raises(APITimeoutError): |
684 | | - self.client.post( |
685 | | - "/v1/inference/chat-completion", |
686 | | - body=cast( |
687 | | - object, |
688 | | - maybe_transform( |
689 | | - dict( |
690 | | - messages=[ |
691 | | - { |
692 | | - "content": "string", |
693 | | - "role": "user", |
694 | | - } |
695 | | - ], |
696 | | - model_id="model_id", |
697 | | - ), |
698 | | - InferenceChatCompletionParamsNonStreaming, |
699 | | - ), |
700 | | - ), |
701 | | - cast_to=httpx.Response, |
702 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
703 | | - ) |
| 683 | + client.inference.with_streaming_response.chat_completion( |
| 684 | + messages=[ |
| 685 | + { |
| 686 | + "content": "string", |
| 687 | + "role": "user", |
| 688 | + } |
| 689 | + ], |
| 690 | + model_id="model_id", |
| 691 | + ).__enter__() |
704 | 692 |
|
705 | 693 | assert _get_open_connections(self.client) == 0 |
706 | 694 |
|
707 | 695 | @mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
708 | 696 | @pytest.mark.respx(base_url=base_url) |
709 | | - def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 697 | + def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: LlamaStackClient) -> None: |
710 | 698 | respx_mock.post("/v1/inference/chat-completion").mock(return_value=httpx.Response(500)) |
711 | 699 |
|
712 | 700 | with pytest.raises(APIStatusError): |
713 | | - self.client.post( |
714 | | - "/v1/inference/chat-completion", |
715 | | - body=cast( |
716 | | - object, |
717 | | - maybe_transform( |
718 | | - dict( |
719 | | - messages=[ |
720 | | - { |
721 | | - "content": "string", |
722 | | - "role": "user", |
723 | | - } |
724 | | - ], |
725 | | - model_id="model_id", |
726 | | - ), |
727 | | - InferenceChatCompletionParamsNonStreaming, |
728 | | - ), |
729 | | - ), |
730 | | - cast_to=httpx.Response, |
731 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
732 | | - ) |
733 | | - |
| 701 | + client.inference.with_streaming_response.chat_completion( |
| 702 | + messages=[ |
| 703 | + { |
| 704 | + "content": "string", |
| 705 | + "role": "user", |
| 706 | + } |
| 707 | + ], |
| 708 | + model_id="model_id", |
| 709 | + ).__enter__() |
734 | 710 | assert _get_open_connections(self.client) == 0 |
735 | 711 |
|
736 | 712 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
@@ -836,6 +812,28 @@ def retry_handler(_request: httpx.Request) -> httpx.Response: |
836 | 812 |
|
837 | 813 | assert response.http_request.headers.get("x-stainless-retry-count") == "42" |
838 | 814 |
|
| 815 | + def test_proxy_environment_variables(self, monkeypatch: pytest.MonkeyPatch) -> None: |
| 816 | + # Test that the proxy environment variables are set correctly |
| 817 | + monkeypatch.setenv("HTTPS_PROXY", "https://example.org") |
| 818 | + |
| 819 | + client = DefaultHttpxClient() |
| 820 | + |
| 821 | + mounts = tuple(client._mounts.items()) |
| 822 | + assert len(mounts) == 1 |
| 823 | + assert mounts[0][0].pattern == "https://" |
| 824 | + |
| 825 | + @pytest.mark.filterwarnings("ignore:.*deprecated.*:DeprecationWarning") |
| 826 | + def test_default_client_creation(self) -> None: |
| 827 | + # Ensure that the client can be initialized without any exceptions |
| 828 | + DefaultHttpxClient( |
| 829 | + verify=True, |
| 830 | + cert=None, |
| 831 | + trust_env=True, |
| 832 | + http1=True, |
| 833 | + http2=False, |
| 834 | + limits=httpx.Limits(max_connections=100, max_keepalive_connections=20), |
| 835 | + ) |
| 836 | + |
839 | 837 | @pytest.mark.respx(base_url=base_url) |
840 | 838 | def test_follow_redirects(self, respx_mock: MockRouter) -> None: |
841 | 839 | # Test that the default follow_redirects=True allows following redirects |
@@ -1495,60 +1493,41 @@ async def test_parse_retry_after_header(self, remaining_retries: int, retry_afte |
1495 | 1493 |
|
1496 | 1494 | @mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1497 | 1495 | @pytest.mark.respx(base_url=base_url) |
1498 | | - async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1496 | + async def test_retrying_timeout_errors_doesnt_leak( |
| 1497 | + self, respx_mock: MockRouter, async_client: AsyncLlamaStackClient |
| 1498 | + ) -> None: |
1499 | 1499 | respx_mock.post("/v1/inference/chat-completion").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
1500 | 1500 |
|
1501 | 1501 | with pytest.raises(APITimeoutError): |
1502 | | - await self.client.post( |
1503 | | - "/v1/inference/chat-completion", |
1504 | | - body=cast( |
1505 | | - object, |
1506 | | - maybe_transform( |
1507 | | - dict( |
1508 | | - messages=[ |
1509 | | - { |
1510 | | - "content": "string", |
1511 | | - "role": "user", |
1512 | | - } |
1513 | | - ], |
1514 | | - model_id="model_id", |
1515 | | - ), |
1516 | | - InferenceChatCompletionParamsNonStreaming, |
1517 | | - ), |
1518 | | - ), |
1519 | | - cast_to=httpx.Response, |
1520 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1521 | | - ) |
| 1502 | + await async_client.inference.with_streaming_response.chat_completion( |
| 1503 | + messages=[ |
| 1504 | + { |
| 1505 | + "content": "string", |
| 1506 | + "role": "user", |
| 1507 | + } |
| 1508 | + ], |
| 1509 | + model_id="model_id", |
| 1510 | + ).__aenter__() |
1522 | 1511 |
|
1523 | 1512 | assert _get_open_connections(self.client) == 0 |
1524 | 1513 |
|
1525 | 1514 | @mock.patch("llama_stack_client._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1526 | 1515 | @pytest.mark.respx(base_url=base_url) |
1527 | | - async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1516 | + async def test_retrying_status_errors_doesnt_leak( |
| 1517 | + self, respx_mock: MockRouter, async_client: AsyncLlamaStackClient |
| 1518 | + ) -> None: |
1528 | 1519 | respx_mock.post("/v1/inference/chat-completion").mock(return_value=httpx.Response(500)) |
1529 | 1520 |
|
1530 | 1521 | with pytest.raises(APIStatusError): |
1531 | | - await self.client.post( |
1532 | | - "/v1/inference/chat-completion", |
1533 | | - body=cast( |
1534 | | - object, |
1535 | | - maybe_transform( |
1536 | | - dict( |
1537 | | - messages=[ |
1538 | | - { |
1539 | | - "content": "string", |
1540 | | - "role": "user", |
1541 | | - } |
1542 | | - ], |
1543 | | - model_id="model_id", |
1544 | | - ), |
1545 | | - InferenceChatCompletionParamsNonStreaming, |
1546 | | - ), |
1547 | | - ), |
1548 | | - cast_to=httpx.Response, |
1549 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1550 | | - ) |
1551 | | - |
| 1522 | + await async_client.inference.with_streaming_response.chat_completion( |
| 1523 | + messages=[ |
| 1524 | + { |
| 1525 | + "content": "string", |
| 1526 | + "role": "user", |
| 1527 | + } |
| 1528 | + ], |
| 1529 | + model_id="model_id", |
| 1530 | + ).__aenter__() |
1552 | 1531 | assert _get_open_connections(self.client) == 0 |
1553 | 1532 |
|
1554 | 1533 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
@@ -1702,6 +1681,28 @@ async def test_main() -> None: |
1702 | 1681 |
|
1703 | 1682 | time.sleep(0.1) |
1704 | 1683 |
|
| 1684 | + async def test_proxy_environment_variables(self, monkeypatch: pytest.MonkeyPatch) -> None: |
| 1685 | + # Test that the proxy environment variables are set correctly |
| 1686 | + monkeypatch.setenv("HTTPS_PROXY", "https://example.org") |
| 1687 | + |
| 1688 | + client = DefaultAsyncHttpxClient() |
| 1689 | + |
| 1690 | + mounts = tuple(client._mounts.items()) |
| 1691 | + assert len(mounts) == 1 |
| 1692 | + assert mounts[0][0].pattern == "https://" |
| 1693 | + |
| 1694 | + @pytest.mark.filterwarnings("ignore:.*deprecated.*:DeprecationWarning") |
| 1695 | + async def test_default_client_creation(self) -> None: |
| 1696 | + # Ensure that the client can be initialized without any exceptions |
| 1697 | + DefaultAsyncHttpxClient( |
| 1698 | + verify=True, |
| 1699 | + cert=None, |
| 1700 | + trust_env=True, |
| 1701 | + http1=True, |
| 1702 | + http2=False, |
| 1703 | + limits=httpx.Limits(max_connections=100, max_keepalive_connections=20), |
| 1704 | + ) |
| 1705 | + |
1705 | 1706 | @pytest.mark.respx(base_url=base_url) |
1706 | 1707 | async def test_follow_redirects(self, respx_mock: MockRouter) -> None: |
1707 | 1708 | # Test that the default follow_redirects=True allows following redirects |
|
0 commit comments