codekingpro/portable-devtools

1import pytest
2from unittest.mock import patch
3from chromadb import CloudClient
4from chromadb.errors import ChromaAuthError, NotFoundError
5from chromadb.auth import UserIdentity
6from chromadb.types import Tenant, Database
7from uuid import uuid4
8
9
10def test_valid_key() -> None:
11    with patch(
12        "chromadb.api.fastapi.FastAPI.get_user_identity"
13    ) as mock_get_user_identity, patch(
14        "chromadb.api.client.AdminClient.get_tenant"
15    ) as mock_get_tenant, patch(
16        "chromadb.api.client.AdminClient.get_database"
17    ) as mock_get_database, patch(
18        "chromadb.api.fastapi.FastAPI.heartbeat"
19    ) as mock_heartbeat:
20        mock_get_user_identity.return_value = UserIdentity(
21            user_id="test_user", tenant="default_tenant", databases=["testdb"]
22        )
23        mock_get_tenant.return_value = Tenant(name="default_tenant")
24        mock_get_database.return_value = Database(
25            id=uuid4(), name="testdb", tenant="default_tenant"
26        )
27        mock_heartbeat.return_value = 1234567890
28
29        client = CloudClient(database="testdb", api_key="valid_token")
30
31        assert client.get_user_identity().user_id == "test_user"
32        assert client.get_user_identity().tenant == "default_tenant"
33        assert client.get_user_identity().databases == ["testdb"]
34
35        settings = client.get_settings()
36        assert settings.chroma_client_auth_credentials == "valid_token"
37        assert (
38            settings.chroma_client_auth_provider
39            == "chromadb.auth.token_authn.TokenAuthClientProvider"
40        )
41
42        assert client.heartbeat() == 1234567890
43
44
45def test_invalid_key() -> None:
46    with patch(
47        "chromadb.api.fastapi.FastAPI.get_user_identity"
48    ) as mock_get_user_identity:
49        mock_get_user_identity.side_effect = ChromaAuthError("Authentication failed")
50
51        with pytest.raises(ChromaAuthError):
52            CloudClient(database="testdb", api_key="invalid_token")
53
54
55# Scoped API key to 1 database tests
56def test_scoped_api_key_to_single_db_with_api_key_only() -> None:
57    with patch(
58        "chromadb.api.fastapi.FastAPI.get_user_identity"
59    ) as mock_get_user_identity, patch(
60        "chromadb.api.client.AdminClient.get_tenant"
61    ) as mock_get_tenant, patch(
62        "chromadb.api.client.AdminClient.get_database"
63    ) as mock_get_database:
64        # mock single db scoped api key
65        mock_get_user_identity.return_value = UserIdentity(
66            user_id="test_user", tenant="123-456-789", databases=["right-db"]
67        )
68        mock_get_tenant.return_value = Tenant(name="123-456-789")
69        mock_get_database.return_value = Database(
70            id=uuid4(), name="right-db", tenant="123-456-789"
71        )
72
73        client = CloudClient(api_key="valid_token")
74
75        # should resolve to single db
76        assert client.database == "right-db"
77        assert client.tenant == "123-456-789"
78
79
80def test_scoped_api_key_to_single_db_with_correct_tenant() -> None:
81    with patch(
82        "chromadb.api.fastapi.FastAPI.get_user_identity"
83    ) as mock_get_user_identity, patch(
84        "chromadb.api.client.AdminClient.get_tenant"
85    ) as mock_get_tenant, patch(
86        "chromadb.api.client.AdminClient.get_database"
87    ) as mock_get_database:
88        mock_get_user_identity.return_value = UserIdentity(
89            user_id="test_user", tenant="123-456-789", databases=["right-db"]
90        )
91        mock_get_tenant.return_value = Tenant(name="123-456-789")
92        mock_get_database.return_value = Database(
93            id=uuid4(), name="right-db", tenant="123-456-789"
94        )
95
96        client = CloudClient(tenant="123-456-789", api_key="valid_token")
97
98        assert client.tenant == "123-456-789"
99        assert client.database == "right-db"
100
101
102def test_scoped_api_key_to_single_db_with_correct_db() -> None:
103    with patch(
104        "chromadb.api.fastapi.FastAPI.get_user_identity"
105    ) as mock_get_user_identity, patch(
106        "chromadb.api.client.AdminClient.get_tenant"
107    ) as mock_get_tenant, patch(
108        "chromadb.api.client.AdminClient.get_database"
109    ) as mock_get_database:
110        mock_get_user_identity.return_value = UserIdentity(
111            user_id="test_user", tenant="123-456-789", databases=["right-db"]
112        )
113        mock_get_tenant.return_value = Tenant(name="123-456-789")
114        mock_get_database.return_value = Database(
115            id=uuid4(), name="right-db", tenant="123-456-789"
116        )
117
118        client = CloudClient(database="right-db", api_key="valid_token")
119
120        assert client.tenant == "123-456-789"
121        assert client.database == "right-db"
122
123
124def test_scoped_api_key_to_single_db_with_correct_tenant_and_db() -> None:
125    with patch(
126        "chromadb.api.fastapi.FastAPI.get_user_identity"
127    ) as mock_get_user_identity, patch(
128        "chromadb.api.client.AdminClient.get_tenant"
129    ) as mock_get_tenant, patch(
130        "chromadb.api.client.AdminClient.get_database"
131    ) as mock_get_database:
132        mock_get_user_identity.return_value = UserIdentity(
133            user_id="test_user", tenant="123-456-789", databases=["right-db"]
134        )
135        mock_get_tenant.return_value = Tenant(name="123-456-789")
136        mock_get_database.return_value = Database(
137            id=uuid4(), name="right-db", tenant="123-456-789"
138        )
139
140        client = CloudClient(
141            tenant="123-456-789", database="right-db", api_key="valid_token"
142        )
143
144        assert client.tenant == "123-456-789"
145        assert client.database == "right-db"
146
147
148def test_scoped_api_key_to_single_db_with_wrong_tenant() -> None:
149    with patch(
150        "chromadb.api.fastapi.FastAPI.get_user_identity"
151    ) as mock_get_user_identity:
152        mock_get_user_identity.return_value = UserIdentity(
153            user_id="test_user", tenant="123-456-789", databases=["right-db"]
154        )
155
156        with pytest.raises(
157            ChromaAuthError,
158            match="Tenant wrong-tenant does not match 123-456-789 from the server. Are you sure the tenant is correct?",
159        ):
160            CloudClient(tenant="wrong-tenant", api_key="valid_token")
161
162
163def test_scoped_api_key_to_single_db_with_wrong_database() -> None:
164    with patch(
165        "chromadb.api.fastapi.FastAPI.get_user_identity"
166    ) as mock_get_user_identity:
167        mock_get_user_identity.return_value = UserIdentity(
168            user_id="test_user", tenant="123-456-789", databases=["right-db"]
169        )
170
171        with pytest.raises(
172            ChromaAuthError,
173            match="Database wrong-db does not match right-db from the server. Are you sure the database is correct?",
174        ):
175            CloudClient(database="wrong-db", api_key="valid_token")
176
177
178def test_scoped_api_key_to_single_db_with_wrong_api_key() -> None:
179    with patch(
180        "chromadb.api.fastapi.FastAPI.get_user_identity"
181    ) as mock_get_user_identity:
182        mock_get_user_identity.side_effect = ChromaAuthError("Permission denied.")
183
184        with pytest.raises(ChromaAuthError, match="Permission denied."):
185            CloudClient(database="right-db", api_key="wrong-api-key")
186
187
188# Scoped API key to multiple databases tests
189def test_scoped_api_key_to_multiple_dbs_with_wrong_tenant() -> None:
190    with patch(
191        "chromadb.api.fastapi.FastAPI.get_user_identity"
192    ) as mock_get_user_identity:
193        mock_get_user_identity.return_value = UserIdentity(
194            user_id="test_user",
195            tenant="123-456-789",
196            databases=["right-db", "another-db"],
197        )
198
199        with pytest.raises(
200            ChromaAuthError,
201            match="Tenant wrong-tenant does not match 123-456-789 from the server. Are you sure the tenant is correct?",
202        ):
203            CloudClient(
204                tenant="wrong-tenant", database="right-db", api_key="valid_token"
205            )
206
207
208def test_scoped_api_key_to_multiple_dbs_with_correct_tenant_and_db() -> None:
209    with patch(
210        "chromadb.api.fastapi.FastAPI.get_user_identity"
211    ) as mock_get_user_identity, patch(
212        "chromadb.api.client.AdminClient.get_tenant"
213    ) as mock_get_tenant, patch(
214        "chromadb.api.client.AdminClient.get_database"
215    ) as mock_get_database:
216        mock_get_user_identity.return_value = UserIdentity(
217            user_id="test_user",
218            tenant="123-456-789",
219            databases=["right-db", "another-db"],
220        )
221        mock_get_tenant.return_value = Tenant(name="123-456-789")
222        mock_get_database.return_value = Database(
223            id=uuid4(), name="right-db", tenant="123-456-789"
224        )
225
226        client = CloudClient(
227            tenant="123-456-789", database="right-db", api_key="valid_token"
228        )
229
230        assert client.tenant == "123-456-789"
231        assert client.database == "right-db"
232
233
234def test_scoped_api_key_to_multiple_dbs_with_nonexistent_database() -> None:
235    with patch(
236        "chromadb.api.fastapi.FastAPI.get_user_identity"
237    ) as mock_get_user_identity, patch(
238        "chromadb.api.client.AdminClient.get_tenant"
239    ) as mock_get_tenant, patch(
240        "chromadb.api.client.AdminClient.get_database"
241    ) as mock_get_database:
242        mock_get_user_identity.return_value = UserIdentity(
243            user_id="test_user",
244            tenant="123-456-789",
245            databases=["right-db", "another-db"],
246        )
247        mock_get_tenant.return_value = Tenant(name="123-456-789")
248        mock_get_database.side_effect = NotFoundError(
249            "Database [wrong-db] not found. Are you sure it exists?"
250        )
251
252        with pytest.raises(
253            NotFoundError,
254            match="Database \\[wrong-db\\] not found. Are you sure it exists?",
255        ):
256            CloudClient(database="wrong-db", api_key="valid_token")
257
258
259def test_scoped_api_key_to_multiple_dbs_with_api_key_only() -> None:
260    with patch(
261        "chromadb.api.fastapi.FastAPI.get_user_identity"
262    ) as mock_get_user_identity:
263        mock_get_user_identity.return_value = UserIdentity(
264            user_id="test_user",
265            tenant="123-456-789",
266            databases=["right-db", "another-db"],
267        )
268
269        with pytest.raises(
270            ChromaAuthError,
271            match="Could not determine a database name from the current authentication method. Please provide a database name.",
272        ):
273            CloudClient(api_key="valid_token")
274
275
276# Unscoped API key tests
277def test_api_key_with_unscoped_tenant() -> None:
278    with patch(
279        "chromadb.api.fastapi.FastAPI.get_user_identity"
280    ) as mock_get_user_identity:
281        mock_get_user_identity.return_value = UserIdentity(
282            user_id="test_user", tenant="*", databases=["right-db"]
283        )
284
285        with pytest.raises(
286            ChromaAuthError,
287            match="Could not determine a tenant from the current authentication method. Please provide a tenant.",
288        ):
289            CloudClient(api_key="valid_token")
290
291
292def test_api_key_with_unscoped_db() -> None:
293    with patch(
294        "chromadb.api.fastapi.FastAPI.get_user_identity"
295    ) as mock_get_user_identity:
296        mock_get_user_identity.return_value = UserIdentity(
297            user_id="test_user", tenant="123-456-789", databases=["*"]
298        )
299
300        with pytest.raises(
301            ChromaAuthError,
302            match="Could not determine a database name from the current authentication method. Please provide a database name.",
303        ):
304            CloudClient(api_key="valid_token")
305
306
307def test_api_key_with_no_db_access() -> None:
308    with patch(
309        "chromadb.api.fastapi.FastAPI.get_user_identity"
310    ) as mock_get_user_identity:
311        mock_get_user_identity.return_value = UserIdentity(
312            user_id="test_user", tenant="123-456-789", databases=[]
313        )
314
315        with pytest.raises(
316            ChromaAuthError,
317            match="Could not determine a database name from the current authentication method. Please provide a database name.",
318        ):
319            CloudClient(api_key="valid_token")
320
321
322def test_api_key_with_no_tenant_access() -> None:
323    with patch(
324        "chromadb.api.fastapi.FastAPI.get_user_identity"
325    ) as mock_get_user_identity:
326        mock_get_user_identity.return_value = UserIdentity(
327            user_id="test_user", tenant=None, databases=["right-db"]
328        )
329
330        with pytest.raises(
331            ChromaAuthError,
332            match="Could not determine a tenant from the current authentication method. Please provide a tenant.",
333        ):
334            CloudClient(api_key="valid_token")
335