codekingpro/portable-devtools

1import asyncio
2from typing import Any, Callable, Generator, cast, Dict, Tuple
3from unittest.mock import MagicMock, patch
4import chromadb
5from chromadb.config import Settings, System
6from chromadb.api import ClientAPI
7import chromadb.server.fastapi
8from chromadb.api.fastapi import FastAPI
9import pytest
10import tempfile
11import os
12
13
14@pytest.fixture
15def ephemeral_api() -> Generator[ClientAPI, None, None]:
16    if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"):
17        pytest.skip("Integration test only")
18    client = chromadb.EphemeralClient()
19    yield client
20    client.clear_system_cache()
21
22
23@pytest.fixture
24def persistent_api() -> Generator[ClientAPI, None, None]:
25    if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"):
26        pytest.skip("Integration test only")
27    client = chromadb.PersistentClient(
28        path=tempfile.gettempdir() + "/test_server",
29    )
30    yield client
31    client.clear_system_cache()
32
33
34HttpAPIFactory = Callable[..., ClientAPI]
35
36
37@pytest.fixture(params=["sync_client", "async_client"])
38def http_api_factory(
39    request: pytest.FixtureRequest,
40) -> Generator[HttpAPIFactory, None, None]:
41    if request.param == "sync_client":
42        with patch("chromadb.api.client.Client._validate_tenant_database"):
43            with patch("chromadb.api.client.Client.get_user_identity"):
44                yield chromadb.HttpClient
45    else:
46        with patch("chromadb.api.async_client.AsyncClient._validate_tenant_database"):
47            with patch("chromadb.api.async_client.AsyncClient.get_user_identity"):
48
49                def factory(*args: Any, **kwargs: Any) -> Any:
50                    cls = asyncio.get_event_loop().run_until_complete(
51                        chromadb.AsyncHttpClient(*args, **kwargs)
52                    )
53                    return cls
54
55                yield cast(HttpAPIFactory, factory)
56
57
58@pytest.fixture()
59def http_api(http_api_factory: HttpAPIFactory) -> Generator[ClientAPI, None, None]:
60    if os.environ.get("CHROMA_SERVER_HTTP_PORT") is not None:
61        port = int(os.environ.get("CHROMA_SERVER_HTTP_PORT"))  # type: ignore
62        client = http_api_factory(port=port)
63    else:
64        client = http_api_factory()
65    yield client
66    client.clear_system_cache()
67
68
69def test_ephemeral_client(ephemeral_api: ClientAPI) -> None:
70    settings = ephemeral_api.get_settings()
71    assert settings.is_persistent is False
72
73
74def test_persistent_client(persistent_api: ClientAPI) -> None:
75    settings = persistent_api.get_settings()
76    assert settings.is_persistent is True
77
78
79def test_http_client(http_api: ClientAPI) -> None:
80    settings = http_api.get_settings()
81    assert (
82        settings.chroma_api_impl == "chromadb.api.fastapi.FastAPI"
83        or settings.chroma_api_impl == "chromadb.api.async_fastapi.AsyncFastAPI"
84    )
85
86
87def test_http_client_with_inconsistent_host_settings(
88    http_api_factory: HttpAPIFactory,
89) -> None:
90    try:
91        http_api_factory(settings=Settings(chroma_server_host="127.0.0.1"))
92    except ValueError as e:
93        assert (
94            str(e)
95            == "Chroma server host provided in settings[127.0.0.1] is different to the one provided in HttpClient: [localhost]"
96        )
97
98
99def test_http_client_with_inconsistent_port_settings(
100    http_api_factory: HttpAPIFactory,
101) -> None:
102    try:
103        http_api_factory(
104            port=8002,
105            settings=Settings(
106                chroma_server_http_port=8001,
107            ),
108        )
109    except ValueError as e:
110        assert (
111            str(e)
112            == "Chroma server http port provided in settings[8001] is different to the one provided in HttpClient: [8002]"
113        )
114
115
116def make_sync_client_factory() -> Tuple[Callable[..., Any], Dict[str, Any]]:
117    captured: Dict[str, Any] = {}
118
119    # takes any positional args to match httpx.Client
120    def factory(*_: Any, **kwargs: Any) -> Any:
121        captured.update(kwargs)
122        session = MagicMock()
123        session.headers = {}
124        return session
125
126    return factory, captured
127
128
129def test_fastapi_uses_http_limits_from_settings() -> None:
130    settings = Settings(
131        chroma_api_impl="chromadb.api.fastapi.FastAPI",
132        chroma_server_host="localhost",
133        chroma_server_http_port=9000,
134        chroma_server_ssl_verify=True,
135        chroma_http_keepalive_secs=12.5,
136        chroma_http_max_connections=64,
137        chroma_http_max_keepalive_connections=16,
138    )
139    system = System(settings)
140
141    factory, captured = make_sync_client_factory()
142
143    with patch.object(FastAPI, "require", side_effect=[MagicMock(), MagicMock()]):
144        with patch("chromadb.api.fastapi.httpx.Client", side_effect=factory):
145            api = FastAPI(system)
146
147    api.stop()
148    limits = captured["limits"]
149    assert limits.keepalive_expiry == 12.5
150    assert limits.max_connections == 64
151    assert limits.max_keepalive_connections == 16
152    assert captured["timeout"] is None
153    assert captured["verify"] is True
154
155
156def test_persistent_client_close() -> None:
157    """Test that close() properly releases resources in PersistentClient."""
158    if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"):
159        pytest.skip("Integration test only")
160
161    with tempfile.TemporaryDirectory() as tmpdir:
162        # Create a client, add some data, and close it
163        client = chromadb.PersistentClient(path=tmpdir)
164        collection = client.create_collection("test_collection")
165        collection.add(
166            ids=["id1", "id2"],
167            documents=["doc1", "doc2"],
168            metadatas=[{"key": "value1"}, {"key": "value2"}],
169        )
170
171        # Save a reference to the system before close() removes it from the cache
172        system = client._system
173
174        # Close the client
175        client.close()
176
177        # Verify the system is stopped
178        assert system._running is False
179
180        # Create a new client with the same path to verify data was persisted
181        client2 = chromadb.PersistentClient(path=tmpdir)
182        collection2 = client2.get_collection("test_collection")
183        results = collection2.get()
184        assert len(results["ids"]) == 2
185        assert "id1" in results["ids"]
186        assert "id2" in results["ids"]
187
188        client2.close()
189        client.clear_system_cache()
190        client2.clear_system_cache()
191
192
193def test_persistent_client_context_manager() -> None:
194    """Test that PersistentClient works as a context manager."""
195    if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"):
196        pytest.skip("Integration test only")
197
198    with tempfile.TemporaryDirectory() as tmpdir:
199        # Use client as context manager
200        with chromadb.PersistentClient(path=tmpdir) as client:
201            # Save a reference to the system before close() removes it from the cache
202            system = client._system
203            collection = client.create_collection("test_collection")
204            collection.add(
205                ids=["id1", "id2"],
206                documents=["doc1", "doc2"],
207                metadatas=[{"key": "value1"}, {"key": "value2"}],
208            )
209
210        # Verify the system is stopped after context exit
211        assert system._running is False
212
213        # Verify data was persisted
214        with chromadb.PersistentClient(path=tmpdir) as client2:
215            collection2 = client2.get_collection("test_collection")
216            results = collection2.get()
217            assert len(results["ids"]) == 2
218
219        client.clear_system_cache()
220        client2.clear_system_cache()
221
222
223def test_ephemeral_client_close() -> None:
224    """Test that close() works with EphemeralClient."""
225    if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"):
226        pytest.skip("Integration test only")
227
228    client = chromadb.EphemeralClient()
229    # Save a reference to the system before close() removes it from the cache
230    system = client._system
231    collection = client.create_collection("test_collection")
232    collection.add(ids=["id1"], documents=["doc1"])
233
234    # Close the client
235    client.close()
236
237    # Verify the system is stopped
238    assert system._running is False
239
240    client.clear_system_cache()
241
242
243def test_ephemeral_client_context_manager() -> None:
244    """Test that EphemeralClient works as a context manager."""
245    if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"):
246        pytest.skip("Integration test only")
247
248    with chromadb.EphemeralClient() as client:
249        # Save a reference to the system before close() removes it from the cache
250        system = client._system
251        collection = client.create_collection("test_collection")
252        collection.add(ids=["id1"], documents=["doc1"])
253        assert system._running is True
254
255    # Verify the system is stopped after context exit
256    assert system._running is False
257
258    client.clear_system_cache()
259
260
261def test_client_close_idempotent() -> None:
262    """Test that calling close() multiple times is a safe no-op."""
263    if os.environ.get("CHROMA_INTEGRATION_TEST_ONLY"):
264        pytest.skip("Integration test only")
265
266    with tempfile.TemporaryDirectory() as tmpdir:
267        client = chromadb.PersistentClient(path=tmpdir)
268        collection = client.create_collection("test_collection")
269        collection.add(ids=["id1"], documents=["doc1"])
270
271        # First close should work normally
272        client.close()
273
274        # Second close should be a no-op, not raise KeyError
275        client.close()
276
277        # Third close should also be safe
278        client.close()
279
280        client.clear_system_cache()
281