|
23 | 23 |
|
24 | 24 | from supermemory import Supermemory, AsyncSupermemory, APIResponseValidationError |
25 | 25 | from supermemory._types import Omit |
26 | | -from supermemory._utils import maybe_transform |
27 | 26 | from supermemory._models import BaseModel, FinalRequestOptions |
28 | | -from supermemory._constants import RAW_RESPONSE_HEADER |
29 | 27 | from supermemory._exceptions import APIStatusError, APITimeoutError, SupermemoryError, APIResponseValidationError |
30 | 28 | from supermemory._base_client import ( |
31 | 29 | DEFAULT_TIMEOUT, |
|
35 | 33 | DefaultAsyncHttpxClient, |
36 | 34 | make_request_options, |
37 | 35 | ) |
38 | | -from supermemory.types.memory_add_params import MemoryAddParams |
39 | 36 |
|
40 | 37 | from .utils import update_env |
41 | 38 |
|
@@ -725,42 +722,25 @@ def test_parse_retry_after_header(self, remaining_retries: int, retry_after: str |
725 | 722 |
|
726 | 723 | @mock.patch("supermemory._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
727 | 724 | @pytest.mark.respx(base_url=base_url) |
728 | | - def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 725 | + def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter, client: Supermemory) -> None: |
729 | 726 | respx_mock.post("/v3/memories").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
730 | 727 |
|
731 | 728 | with pytest.raises(APITimeoutError): |
732 | | - self.client.post( |
733 | | - "/v3/memories", |
734 | | - body=cast( |
735 | | - object, |
736 | | - maybe_transform( |
737 | | - dict(content="This is a detailed article about machine learning concepts..."), MemoryAddParams |
738 | | - ), |
739 | | - ), |
740 | | - cast_to=httpx.Response, |
741 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
742 | | - ) |
| 729 | + client.memories.with_streaming_response.add( |
| 730 | + content="This is a detailed article about machine learning concepts..." |
| 731 | + ).__enter__() |
743 | 732 |
|
744 | 733 | assert _get_open_connections(self.client) == 0 |
745 | 734 |
|
746 | 735 | @mock.patch("supermemory._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
747 | 736 | @pytest.mark.respx(base_url=base_url) |
748 | | - def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 737 | + def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter, client: Supermemory) -> None: |
749 | 738 | respx_mock.post("/v3/memories").mock(return_value=httpx.Response(500)) |
750 | 739 |
|
751 | 740 | with pytest.raises(APIStatusError): |
752 | | - self.client.post( |
753 | | - "/v3/memories", |
754 | | - body=cast( |
755 | | - object, |
756 | | - maybe_transform( |
757 | | - dict(content="This is a detailed article about machine learning concepts..."), MemoryAddParams |
758 | | - ), |
759 | | - ), |
760 | | - cast_to=httpx.Response, |
761 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
762 | | - ) |
763 | | - |
| 741 | + client.memories.with_streaming_response.add( |
| 742 | + content="This is a detailed article about machine learning concepts..." |
| 743 | + ).__enter__() |
764 | 744 | assert _get_open_connections(self.client) == 0 |
765 | 745 |
|
766 | 746 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
@@ -1568,42 +1548,29 @@ async def test_parse_retry_after_header(self, remaining_retries: int, retry_afte |
1568 | 1548 |
|
1569 | 1549 | @mock.patch("supermemory._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1570 | 1550 | @pytest.mark.respx(base_url=base_url) |
1571 | | - async def test_retrying_timeout_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1551 | + async def test_retrying_timeout_errors_doesnt_leak( |
| 1552 | + self, respx_mock: MockRouter, async_client: AsyncSupermemory |
| 1553 | + ) -> None: |
1572 | 1554 | respx_mock.post("/v3/memories").mock(side_effect=httpx.TimeoutException("Test timeout error")) |
1573 | 1555 |
|
1574 | 1556 | with pytest.raises(APITimeoutError): |
1575 | | - await self.client.post( |
1576 | | - "/v3/memories", |
1577 | | - body=cast( |
1578 | | - object, |
1579 | | - maybe_transform( |
1580 | | - dict(content="This is a detailed article about machine learning concepts..."), MemoryAddParams |
1581 | | - ), |
1582 | | - ), |
1583 | | - cast_to=httpx.Response, |
1584 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1585 | | - ) |
| 1557 | + await async_client.memories.with_streaming_response.add( |
| 1558 | + content="This is a detailed article about machine learning concepts..." |
| 1559 | + ).__aenter__() |
1586 | 1560 |
|
1587 | 1561 | assert _get_open_connections(self.client) == 0 |
1588 | 1562 |
|
1589 | 1563 | @mock.patch("supermemory._base_client.BaseClient._calculate_retry_timeout", _low_retry_timeout) |
1590 | 1564 | @pytest.mark.respx(base_url=base_url) |
1591 | | - async def test_retrying_status_errors_doesnt_leak(self, respx_mock: MockRouter) -> None: |
| 1565 | + async def test_retrying_status_errors_doesnt_leak( |
| 1566 | + self, respx_mock: MockRouter, async_client: AsyncSupermemory |
| 1567 | + ) -> None: |
1592 | 1568 | respx_mock.post("/v3/memories").mock(return_value=httpx.Response(500)) |
1593 | 1569 |
|
1594 | 1570 | with pytest.raises(APIStatusError): |
1595 | | - await self.client.post( |
1596 | | - "/v3/memories", |
1597 | | - body=cast( |
1598 | | - object, |
1599 | | - maybe_transform( |
1600 | | - dict(content="This is a detailed article about machine learning concepts..."), MemoryAddParams |
1601 | | - ), |
1602 | | - ), |
1603 | | - cast_to=httpx.Response, |
1604 | | - options={"headers": {RAW_RESPONSE_HEADER: "stream"}}, |
1605 | | - ) |
1606 | | - |
| 1571 | + await async_client.memories.with_streaming_response.add( |
| 1572 | + content="This is a detailed article about machine learning concepts..." |
| 1573 | + ).__aenter__() |
1607 | 1574 | assert _get_open_connections(self.client) == 0 |
1608 | 1575 |
|
1609 | 1576 | @pytest.mark.parametrize("failures_before_success", [0, 2, 4]) |
|
0 commit comments