codekingpro/portable-devtools

1import unittest
2import os
3from unittest.mock import patch, Mock
4import pytest
5import chromadb
6import chromadb.config
7from chromadb.db.system import SysDB
8from chromadb.ingest import Consumer, Producer
9
10
11class GetDBTest(unittest.TestCase):
12    @patch("chromadb.db.impl.sqlite.SqliteDB", autospec=True)
13    def test_default_db(self, mock: Mock) -> None:
14        system = chromadb.config.System(
15            chromadb.config.Settings(persist_directory="./foo")
16        )
17        system.instance(SysDB)
18        assert mock.called
19
20    @patch("chromadb.db.impl.sqlite.SqliteDB", autospec=True)
21    def test_sqlite_sysdb(self, mock: Mock) -> None:
22        system = chromadb.config.System(
23            chromadb.config.Settings(
24                chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB",
25                persist_directory="./foo",
26            )
27        )
28        system.instance(SysDB)
29        assert mock.called
30
31    @patch("chromadb.db.impl.sqlite.SqliteDB", autospec=True)
32    def test_sqlite_queue(self, mock: Mock) -> None:
33        system = chromadb.config.System(
34            chromadb.config.Settings(
35                chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB",
36                chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB",
37                chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB",
38                persist_directory="./foo",
39            )
40        )
41        system.instance(Producer)
42        system.instance(Consumer)
43        assert mock.called
44
45
46class GetAPITest(unittest.TestCase):
47    @patch("chromadb.api.segment.SegmentAPI", autospec=True)
48    @patch.dict(
49        os.environ, {"CHROMA_API_IMPL": "chromadb.api.segment.SegmentAPI"}, clear=True
50    )
51    def test_local(self, mock_api: Mock) -> None:
52        client = chromadb.Client(chromadb.config.Settings(persist_directory="./foo"))
53        assert mock_api.called
54        client.clear_system_cache()
55
56    @patch("chromadb.db.impl.sqlite.SqliteDB", autospec=True)
57    @patch.dict(
58        os.environ, {"CHROMA_API_IMPL": "chromadb.api.segment.SegmentAPI"}, clear=True
59    )
60    def test_local_db(self, mock_db: Mock) -> None:
61        client = chromadb.Client(chromadb.config.Settings(persist_directory="./foo"))
62        assert mock_db.called
63        client.clear_system_cache()
64
65    @patch("chromadb.api.fastapi.FastAPI", autospec=True)
66    @patch.dict(os.environ, {}, clear=True)
67    def test_fastapi(self, mock: Mock) -> None:
68        client = chromadb.Client(
69            chromadb.config.Settings(
70                chroma_api_impl="chromadb.api.fastapi.FastAPI",
71                persist_directory="./foo",
72                chroma_server_host="foo",
73                chroma_server_http_port=80,
74            )
75        )
76        assert mock.called
77        client.clear_system_cache()
78
79    @patch("chromadb.api.fastapi.FastAPI", autospec=True)
80    @patch.dict(os.environ, {}, clear=True)
81    def test_settings_pass_to_fastapi(self, mock: Mock) -> None:
82        settings = chromadb.config.Settings(
83            chroma_api_impl="chromadb.api.fastapi.FastAPI",
84            chroma_server_host="foo",
85            chroma_server_http_port=80,
86            chroma_server_headers={"foo": "bar"},
87        )
88        client = chromadb.Client(settings)
89
90        # Check that the mock was called
91        assert mock.called
92
93        # Retrieve the arguments with which the mock was called
94        # `call_args` returns a tuple, where the first element is a tuple of positional arguments
95        # and the second element is a dictionary of keyword arguments. We assume here that
96        # the settings object is passed as a positional argument.
97        args, kwargs = mock.call_args
98        passed_settings = args[0] if args else None
99
100        # Check if the settings passed to the mock match the settings we used
101        # raise Exception(passed_settings.settings)
102        assert passed_settings.settings == settings
103        client.clear_system_cache()
104
105
106def test_legacy_values() -> None:
107    with pytest.raises(ValueError):
108        client = chromadb.Client(
109            chromadb.config.Settings(
110                chroma_api_impl="chromadb.api.local.LocalAPI",
111                persist_directory="./foo",
112                chroma_server_host="foo",
113                chroma_server_http_port=80,
114            )
115        )
116        client.clear_system_cache()
117