chromadb/test/test_chroma.py

import unittest
import os
from unittest.mock import patch, Mock
import pytest
import chromadb
import chromadb.config
from chromadb.db.system import SysDB
from chromadb.ingest import Consumer, Producer


class GetDBTest(unittest.TestCase):
    @patch("chromadb.db.impl.sqlite.SqliteDB", autospec=True)
    def test_default_db(self, mock: Mock) -> None:
        system = chromadb.config.System(
            chromadb.config.Settings(persist_directory="./foo")
        )
        system.instance(SysDB)
        assert mock.called

    @patch("chromadb.db.impl.sqlite.SqliteDB", autospec=True)
    def test_sqlite_sysdb(self, mock: Mock) -> None:
        system = chromadb.config.System(
            chromadb.config.Settings(
                chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB",
                persist_directory="./foo",
            )
        )
        system.instance(SysDB)
        assert mock.called

    @patch("chromadb.db.impl.sqlite.SqliteDB", autospec=True)
    def test_sqlite_queue(self, mock: Mock) -> None:
        system = chromadb.config.System(
            chromadb.config.Settings(
                chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB",
                chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB",
                chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB",
                persist_directory="./foo",
            )
        )
        system.instance(Producer)
        system.instance(Consumer)
        assert mock.called


class GetAPITest(unittest.TestCase):
    @patch("chromadb.api.segment.SegmentAPI", autospec=True)
    @patch.dict(os.environ, {}, clear=True)
    def test_local(self, mock_api: Mock) -> None:
        client = chromadb.Client(chromadb.config.Settings(persist_directory="./foo"))
        assert mock_api.called
        client.clear_system_cache()

    @patch("chromadb.db.impl.sqlite.SqliteDB", autospec=True)
    @patch.dict(os.environ, {}, clear=True)
    def test_local_db(self, mock_db: Mock) -> None:
        client = chromadb.Client(chromadb.config.Settings(persist_directory="./foo"))
        assert mock_db.called
        client.clear_system_cache()

    @patch("chromadb.api.fastapi.FastAPI", autospec=True)
    @patch.dict(os.environ, {}, clear=True)
    def test_fastapi(self, mock: Mock) -> None:
        client = chromadb.Client(
            chromadb.config.Settings(
                chroma_api_impl="chromadb.api.fastapi.FastAPI",
                persist_directory="./foo",
                chroma_server_host="foo",
                chroma_server_http_port="80",
            )
        )
        assert mock.called
        client.clear_system_cache()

    @patch("chromadb.api.fastapi.FastAPI", autospec=True)
    @patch.dict(os.environ, {}, clear=True)
    def test_settings_pass_to_fastapi(self, mock: Mock) -> None:
        settings = chromadb.config.Settings(
            chroma_api_impl="chromadb.api.fastapi.FastAPI",
            chroma_server_host="foo",
            chroma_server_http_port="80",
            chroma_server_headers={"foo": "bar"},
        )
        client = chromadb.Client(settings)

        # Check that the mock was called
        assert mock.called

        # Retrieve the arguments with which the mock was called
        # `call_args` returns a tuple, where the first element is a tuple of positional arguments
        # and the second element is a dictionary of keyword arguments. We assume here that
        # the settings object is passed as a positional argument.
        args, kwargs = mock.call_args
        passed_settings = args[0] if args else None

        # Check if the settings passed to the mock match the settings we used
        # raise Exception(passed_settings.settings)
        assert passed_settings.settings == settings
        client.clear_system_cache()


def test_legacy_values() -> None:
    with pytest.raises(ValueError):
        client = chromadb.Client(
            chromadb.config.Settings(
                chroma_api_impl="chromadb.api.local.LocalAPI",
                persist_directory="./foo",
                chroma_server_host="foo",
                chroma_server_http_port="80",
            )
        )
        client.clear_system_cache()