codekingpro/portable-devtools

1import pytest
2from unittest.mock import MagicMock
3from chromadb.api.shared_system_client import SharedSystemClient
4from chromadb.api.base_http_client import BaseHTTPClient
5from chromadb.config import System
6from typing import Optional, Dict, Generator
7
8
9@pytest.fixture(autouse=True)
10def clear_cache() -> Generator[None, None, None]:
11    """Automatically clear the system cache before and after each test."""
12    SharedSystemClient.clear_system_cache()
13    yield
14    SharedSystemClient.clear_system_cache()
15
16
17def create_mock_http_client(
18    api_url: Optional[str] = None,
19    headers: Optional[Dict[str, str]] = None,
20) -> MagicMock:
21    """Create a mock BaseHTTPClient instance with the specified configuration."""
22    mock_server_api = MagicMock(spec=BaseHTTPClient)
23
24    mock_server_api.get_api_url.return_value = api_url or ""
25    mock_server_api.get_request_headers.return_value = headers or {}
26
27    return mock_server_api
28
29
30def register_mock_system(system_id: str, mock_server_api: MagicMock) -> MagicMock:
31    """Register a mock system with the given ID and server API."""
32    mock_system = MagicMock(spec=System)
33    mock_system.instance.return_value = mock_server_api
34    SharedSystemClient._identifier_to_system[system_id] = mock_system
35    return mock_system
36
37
38def test_extracts_api_key_from_chroma_cloud_client() -> None:
39    mock_server_api = create_mock_http_client(
40        api_url="https://api.trychroma.com/api/v2",
41        headers={"X-Chroma-Token": "test-api-key-123"},
42    )
43    register_mock_system("test-id", mock_server_api)
44
45    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
46
47    assert api_key == "test-api-key-123"
48
49
50def test_extracts_api_key_with_lowercase_header() -> None:
51    mock_server_api = create_mock_http_client(
52        api_url="https://api.trychroma.com/api/v2",
53        headers={"x-chroma-token": "test-api-key-456"},
54    )
55    register_mock_system("test-id", mock_server_api)
56
57    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
58
59    assert api_key == "test-api-key-456"
60
61
62def test_extracts_api_key_from_gcp_chroma_cloud_client() -> None:
63    mock_server_api = create_mock_http_client(
64        api_url="https://dummy.gcp.trychroma.com/api/v2",
65        headers={"X-Chroma-Token": "gcp-test-api-key"},
66    )
67    register_mock_system("test-id", mock_server_api)
68
69    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
70
71    assert api_key == "gcp-test-api-key"
72
73
74def test_skips_non_chroma_cloud_clients() -> None:
75    mock_server_api = create_mock_http_client(
76        api_url="https://localhost:8000/api/v2",
77        headers={"X-Chroma-Token": "local-api-key"},
78    )
79    register_mock_system("test-id", mock_server_api)
80
81    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
82
83    assert api_key is None
84
85
86def test_skips_clients_without_api_url() -> None:
87    mock_server_api = create_mock_http_client(
88        api_url=None,
89        headers={"X-Chroma-Token": "test-api-key"},
90    )
91    register_mock_system("test-id", mock_server_api)
92
93    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
94
95    assert api_key is None
96
97
98def test_returns_none_when_no_api_key_in_headers() -> None:
99    mock_server_api = create_mock_http_client(
100        api_url="https://api.trychroma.com/api/v2",
101        headers={},
102    )
103    register_mock_system("test-id", mock_server_api)
104
105    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
106
107    assert api_key is None
108
109
110def test_returns_first_api_key_found_from_multiple_clients() -> None:
111    mock_server_api_1 = create_mock_http_client(
112        api_url="https://api.trychroma.com/api/v2",
113        headers={"X-Chroma-Token": "first-key"},
114    )
115    mock_server_api_2 = create_mock_http_client(
116        api_url="https://api.trychroma.com/api/v2",
117        headers={"X-Chroma-Token": "second-key"},
118    )
119    register_mock_system("test-id-1", mock_server_api_1)
120    register_mock_system("test-id-2", mock_server_api_2)
121
122    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
123
124    assert api_key == "first-key"
125
126
127def test_handles_exception_gracefully() -> None:
128    mock_system = MagicMock(spec=System)
129    mock_system.instance.side_effect = Exception("Test exception")
130    SharedSystemClient._identifier_to_system["test-id"] = mock_system
131
132    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
133
134    assert api_key is None
135
136
137def test_returns_none_when_no_clients_exist() -> None:
138    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
139
140    assert api_key is None
141
142
143def test_skips_non_http_clients() -> None:
144    """Test that non-BaseHTTPClient instances are skipped."""
145    mock_server_api = MagicMock()  # Not a BaseHTTPClient
146    register_mock_system("test-id", mock_server_api)
147
148    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
149
150    assert api_key is None
151
152
153def test_extracts_api_key_with_mixed_case_header() -> None:
154    mock_server_api = create_mock_http_client(
155        api_url="https://api.trychroma.com/api/v2",
156        headers={"X-CHROMA-TOKEN": "mixed-case-key"},
157    )
158    register_mock_system("test-id", mock_server_api)
159
160    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
161
162    assert api_key == "mixed-case-key"
163
164
165def test_multiple_clients_returns_one_key() -> None:
166    """Test that multiple clients return one of the available keys."""
167    mock_api_1 = create_mock_http_client(
168        api_url="https://api.trychroma.com/api/v2",
169        headers={"X-Chroma-Token": "key-1"},
170    )
171    mock_api_2 = create_mock_http_client(
172        api_url="https://api.trychroma.com/api/v2",
173        headers={"X-Chroma-Token": "key-2"},
174    )
175    register_mock_system("id-1", mock_api_1)
176    register_mock_system("id-2", mock_api_2)
177
178    api_key = SharedSystemClient.get_chroma_cloud_api_key_from_clients()
179
180    assert api_key in ["key-1", "key-2"]
181