codekingpro/portable-devtools

1from chromadb.config import Component, System, Settings
2from overrides import overrides
3from threading import local
4from unittest.mock import patch
5import pytest
6import os
7import random
8
9data = local()  # use thread local just in case tests ever run in parallel
10
11
12def reset() -> None:
13    global data
14    data.starts = []
15    data.stops = []
16    data.inits = []
17
18
19class ComponentA(Component):
20    def __init__(self, system: System):
21        data.inits += "A"
22        super().__init__(system)
23        self.require(ComponentB)
24        self.require(ComponentC)
25
26    @overrides
27    def start(self) -> None:
28        data.starts += "A"
29
30    @overrides
31    def stop(self) -> None:
32        data.stops += "A"
33
34
35class ComponentB(Component):
36    def __init__(self, system: System):
37        data.inits += "B"
38        super().__init__(system)
39        self.require(ComponentC)
40        self.require(ComponentD)
41
42    @overrides
43    def start(self) -> None:
44        data.starts += "B"
45
46    @overrides
47    def stop(self) -> None:
48        data.stops += "B"
49
50
51class ComponentC(Component):
52    def __init__(self, system: System):
53        data.inits += "C"
54        super().__init__(system)
55        self.require(ComponentD)
56
57    @overrides
58    def start(self) -> None:
59        data.starts += "C"
60
61    @overrides
62    def stop(self) -> None:
63        data.stops += "C"
64
65
66class ComponentD(Component):
67    def __init__(self, system: System):
68        data.inits += "D"
69        super().__init__(system)
70
71    @overrides
72    def start(self) -> None:
73        data.starts += "D"
74
75    @overrides
76    def stop(self) -> None:
77        data.stops += "D"
78
79
80# Dependency Graph for tests:
81# ┌───┐
82# │ A │
83# └┬─┬┘
84#  │┌▽──┐
85#  ││ B │
86#  │└┬─┬┘
87# ┌▽─▽┐│
88# │ C ││
89# └┬──┘│
90# ┌▽───▽┐
91# │  D  │
92# └─────┘
93
94
95def test_leaf_only() -> None:
96    settings = Settings()
97    system = System(settings)
98
99    reset()
100
101    d = system.instance(ComponentD)
102    assert isinstance(d, ComponentD)
103
104    assert data.inits == ["D"]
105    system.start()
106    assert data.starts == ["D"]
107    system.stop()
108    assert data.stops == ["D"]
109
110
111def test_partial() -> None:
112    settings = Settings()
113    system = System(settings)
114
115    reset()
116
117    c = system.instance(ComponentC)
118    assert isinstance(c, ComponentC)
119
120    assert data.inits == ["C", "D"]
121    system.start()
122    assert data.starts == ["D", "C"]
123    system.stop()
124    assert data.stops == ["C", "D"]
125
126
127def test_system_startup() -> None:
128    settings = Settings()
129    system = System(settings)
130
131    reset()
132
133    a = system.instance(ComponentA)
134    assert isinstance(a, ComponentA)
135
136    assert data.inits == ["A", "B", "C", "D"]
137    system.start()
138    assert data.starts == ["D", "C", "B", "A"]
139    system.stop()
140    assert data.stops == ["A", "B", "C", "D"]
141
142
143def test_system_override_order() -> None:
144    settings = Settings()
145    system = System(settings)
146
147    reset()
148
149    system.instance(ComponentA)
150
151    # Deterministically shuffle the instances map to prove that topsort is actually
152    # working and not just implicitly working because of insertion order.
153
154    # This causes the test to actually fail if the deps are not wired up correctly.
155    random.seed(0)
156    entries = list(system._instances.items())
157    random.shuffle(entries)
158    system._instances = {k: v for k, v in entries}
159
160    system.start()
161    assert data.starts == ["D", "C", "B", "A"]
162    system.stop()
163    assert data.stops == ["A", "B", "C", "D"]
164
165
166class ComponentZ(Component):
167    def __init__(self, system: System):
168        super().__init__(system)
169        self.require(ComponentC)
170
171    @overrides
172    def start(self) -> None:
173        pass
174
175    @overrides
176    def stop(self) -> None:
177        pass
178
179
180def test_runtime_dependencies() -> None:
181    settings = Settings()
182    system = System(settings)
183
184    reset()
185
186    # Nothing to do, no components were requested prior to start
187    system.start()
188    assert data.starts == []
189
190    # Constructs dependencies and starts them in the correct order
191    ComponentZ(system)
192    assert data.starts == ["D", "C"]
193    system.stop()
194    assert data.stops == ["C", "D"]
195
196
197def test_http_client_setting_defaults() -> None:
198    settings = Settings()
199    assert settings.chroma_http_keepalive_secs == 40.0
200    assert settings.chroma_http_max_connections is None
201    assert settings.chroma_http_max_keepalive_connections is None
202
203
204def test_http_client_setting_overrides() -> None:
205    settings = Settings(
206        chroma_http_keepalive_secs=5.5,
207        chroma_http_max_connections=123,
208        chroma_http_max_keepalive_connections=17,
209    )
210    assert settings.chroma_http_keepalive_secs == 5.5
211    assert settings.chroma_http_max_connections == 123
212    assert settings.chroma_http_max_keepalive_connections == 17
213
214
215@patch.dict(os.environ, {"CHROMA_API_IMPL": "my_api_impl"}, clear=True)
216def test_uses_env() -> None:
217    settings = Settings()
218    assert settings.chroma_api_impl == "my_api_impl"
219
220
221@patch.dict(os.environ, {"MY_ENV_VAR": "my_env_var"}, clear=True)
222def test_ignores_extra_env_vars() -> None:
223    settings = Settings()
224    with pytest.raises(AttributeError):
225        _ = settings.my_env_var
226
227
228def test_local_ignores_extra_settings_param() -> None:
229    settings = Settings(extra_param="asdsdsds", tenant_id="test")
230    # does not error if the extra param is present in the settings object
231    assert settings.tenant_id == "test"
232    # but it should error if the extra param is accessed
233    with pytest.raises(AttributeError):
234        _ = settings.extra_param
235