test: Chroma - add unit tests (#3175)

* chroma: add unit tests

* mv test

Author: anakin87

integrations/chroma/tests/test_document_store.py

@@ -9,7 +9,7 @@
 
 import pytest
 from chromadb.api.shared_system_client import SharedSystemClient
-from haystack.dataclasses import ByteStream, Document
+from haystack.dataclasses import ByteStream, Document, SparseEmbedding
 from haystack.testing.document_store import (
     TEST_EMBEDDING_1,
     CountDocumentsByFilterTest,
@@ -27,6 +27,8 @@
 )
 
 from haystack_integrations.document_stores.chroma import ChromaDocumentStore
+from haystack_integrations.document_stores.chroma.errors import ChromaDocumentStoreConfigError
+from haystack_integrations.document_stores.chroma.utils import get_embedding_function
 
 
 @pytest.fixture
@@ -42,6 +44,11 @@ def clear_chroma_system_cache():
     SharedSystemClient.clear_system_cache()
 
 
+def test_get_embedding_function_invalid_name_raises():
+    with pytest.raises(ChromaDocumentStoreConfigError, match="Invalid function name"):
+        get_embedding_function("NonExistentEmbeddingFunction")
+
+
 class TestDocumentStoreUnit:
     def test_init_in_memory(self):
         store = ChromaDocumentStore()
@@ -125,6 +132,92 @@ def test_same_collection_name_reinitialization(self):
         ChromaDocumentStore("test_1")
         ChromaDocumentStore("test_1")
 
+    def test_ensure_initialized_invalid_client_settings_raises(self):
+        with mock.patch(
+            "haystack_integrations.document_stores.chroma.document_store.Settings",
+            side_effect=ValueError("bad setting"),
+        ):
+            store = ChromaDocumentStore(client_settings={"foo": "bar"})
+            with pytest.raises(ValueError, match="Invalid client_settings"):
+                store._ensure_initialized()
+
+    def test_infer_type_from_value_fallback_for_unknown_type(self):
+        assert ChromaDocumentStore._infer_type_from_value(None) == "keyword"
+        assert ChromaDocumentStore._infer_type_from_value(["a", "b"]) == "keyword"
+
+    def test_count_unique_metadata_empty_returns_zero_counts(self):
+        assert ChromaDocumentStore._count_unique_metadata(None, ["a", "b"]) == {"a": 0, "b": 0}
+        assert ChromaDocumentStore._count_unique_metadata([], ["x"]) == {"x": 0}
+
+    def test_compute_field_min_max_skips_non_scalar_values(self):
+        metadatas = [{"cat": ["a", "b"]}, {"cat": "X"}, {"cat": "Z"}]
+        result = ChromaDocumentStore._compute_field_min_max(metadatas, "cat")
+        assert result == {"min": "X", "max": "Z"}
+
+    @pytest.mark.parametrize(
+        "result",
+        [
+            {"ids": ["1"], "documents": None, "metadatas": [{"cat": "A"}]},
+            {"ids": ["1"], "documents": ["hello world"], "metadatas": [{"cat": "A"}]},
+        ],
+        ids=["documents_none", "no_matches"],
+    )
+    def test_compute_field_unique_values_with_search_term_edge_cases(self, result):
+        values, total = ChromaDocumentStore._compute_field_unique_values(result, "cat", "absent", 0, 10)
+        assert values == []
+        assert total == 0
+
+    def test_filter_metadata_discards_unsupported_types(self, caplog):
+        meta = {"ok": "x", "also_ok": None, "bad": {"nested": 1}, "worse": object()}
+        with caplog.at_level(logging.WARNING):
+            result = ChromaDocumentStore._filter_metadata(meta)
+        assert result == {"ok": "x", "also_ok": None}
+        assert "bad" in caplog.text and "worse" in caplog.text
+
+    def test_convert_document_to_chroma_rejects_non_document(self):
+        with pytest.raises(ValueError, match="must contain a list of objects of type Document"):
+            ChromaDocumentStore._convert_document_to_chroma("not a document")  # type: ignore[arg-type]
+
+    def test_convert_document_to_chroma_warns_on_sparse_embedding(self, caplog):
+        doc = Document(content="hello", sparse_embedding=SparseEmbedding(indices=[0, 1], values=[0.1, 0.2]))
+        with caplog.at_level(logging.WARNING):
+            data = ChromaDocumentStore._convert_document_to_chroma(doc)
+        assert data is not None
+        assert "sparse_embedding" in caplog.text
+
+    @pytest.mark.parametrize(
+        ("result", "expected_embedding"),
+        [
+            (
+                {"ids": ["1"], "documents": ["c"], "metadatas": [{"k": "v"}], "embeddings": [[0.1, 0.2]]},
+                [0.1, 0.2],
+            ),
+            ({"ids": ["1"], "documents": ["c"], "metadatas": [None]}, None),
+        ],
+        ids=["list_embeddings", "no_embeddings"],
+    )
+    def test_get_result_to_documents_embedding_variants(self, result, expected_embedding):
+        docs = ChromaDocumentStore._get_result_to_documents(result)  # type: ignore[arg-type]
+        assert docs[0].embedding == expected_embedding
+
+    @pytest.mark.parametrize(
+        ("result", "check"),
+        [
+            ({"documents": None}, lambda docs: docs == []),
+            (
+                {"ids": [["a", "b"]], "documents": [["c1", "c2"]], "metadatas": [[{"k": "v"}]]},
+                lambda docs: docs[0][0].meta == {"k": "v"} and docs[0][1].meta == {},
+            ),
+            (
+                {"ids": [["a"]], "documents": [["c"]], "metadatas": [[{"k": "v"}]]},
+                lambda docs: docs[0][0].embedding is None and docs[0][0].score is None,
+            ),
+        ],
+        ids=["documents_none", "metadata_index_error", "no_embeddings_no_distances"],
+    )
+    def test_query_result_to_documents_edge_cases(self, result, check):
+        assert check(ChromaDocumentStore._query_result_to_documents(result))  # type: ignore[arg-type]
+
 
 @pytest.mark.integration
 class TestDocumentStore(

integrations/chroma/tests/test_document_store_async.py

@@ -14,6 +14,23 @@
 from haystack_integrations.document_stores.chroma import ChromaDocumentStore
 
 
+@pytest.mark.asyncio
+class TestDocumentStoreAsyncUnit:
+    async def test_ensure_initialized_async_requires_host_and_port(self):
+        store = ChromaDocumentStore()
+        with pytest.raises(ValueError, match="Async support"):
+            await store._ensure_initialized_async()
+
+    async def test_ensure_initialized_async_invalid_client_settings_raises(self):
+        with mock.patch(
+            "haystack_integrations.document_stores.chroma.document_store.Settings",
+            side_effect=ValueError("bad setting"),
+        ):
+            store = ChromaDocumentStore(host="localhost", port=8000, client_settings={"foo": "bar"})
+            with pytest.raises(ValueError, match="Invalid client_settings"):
+                await store._ensure_initialized_async()
+
+
 @pytest.mark.skipif(
     sys.platform == "win32",
     reason="We do not run the Chroma server on Windows and async is only supported with HTTP connections",

integrations/chroma/tests/test_filters.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from haystack_integrations.document_stores.chroma.errors import ChromaDocumentStoreFilterError
+from haystack_integrations.document_stores.chroma.filters import (
+    _convert_filters,
+    _parse_comparison_condition,
+    _parse_logical_condition,
+)
+
+
+def test_id_filter_with_empty_value_raises():
+    with pytest.raises(ChromaDocumentStoreFilterError, match="id filter only supports"):
+        _convert_filters({"field": "id", "operator": "==", "value": ""})
+
+
+@pytest.mark.parametrize(
+    ("condition", "match"),
+    [
+        ({"conditions": []}, "'operator' key missing"),
+        ({"operator": "AND"}, "'conditions' key missing"),
+        ({"operator": "XOR", "conditions": []}, "Unknown operator"),
+    ],
+    ids=["missing_operator", "missing_conditions", "unknown_operator"],
+)
+def test_parse_logical_condition_errors(condition, match):
+    with pytest.raises(ChromaDocumentStoreFilterError, match=match):
+        _parse_logical_condition(condition)
+
+
+@pytest.mark.parametrize(
+    ("condition", "match"),
+    [
+        ({"operator": "==", "value": "x"}, "'field' key missing"),
+        ({"field": "meta.a", "value": "x"}, "'operator' key missing"),
+        ({"field": "meta.a", "operator": "=="}, "'value' key missing"),
+        ({"field": "meta.a", "operator": "~~", "value": "x"}, "Unknown operator"),
+    ],
+    ids=["missing_field", "missing_operator", "missing_value", "unknown_operator"],
+)
+def test_parse_comparison_condition_errors(condition, match):
+    with pytest.raises(ChromaDocumentStoreFilterError, match=match):
+        _parse_comparison_condition(condition)

integrations/chroma/tests/test_retriever.py

@@ -2,7 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from unittest import mock
+
 import pytest
+from haystack.dataclasses import Document
 from haystack.document_stores.types import FilterPolicy
 
 from haystack_integrations.components.retrievers.chroma import ChromaEmbeddingRetriever, ChromaQueryTextRetriever
@@ -102,6 +105,29 @@ def test_from_dict_no_filter_policy(self, request):
         assert retriever.top_k == 42
         assert retriever.filter_policy == FilterPolicy.REPLACE  # default even if not specified
 
+    def test_run_delegates_to_document_store_search(self):
+        ds = mock.Mock(spec=ChromaDocumentStore)
+        expected = [Document(content="hit")]
+        ds.search.return_value = [expected]
+        retriever = ChromaQueryTextRetriever(ds, top_k=5)
+
+        result = retriever.run(query="q")
+
+        ds.search.assert_called_once_with(["q"], 5, {})
+        assert result == {"documents": expected}
+
+    @pytest.mark.asyncio
+    async def test_run_async_delegates_to_document_store_search_async(self):
+        ds = mock.Mock(spec=ChromaDocumentStore)
+        expected = [Document(content="hit")]
+        ds.search_async = mock.AsyncMock(return_value=[expected])
+        retriever = ChromaQueryTextRetriever(ds, top_k=3)
+
+        result = await retriever.run_async(query="q")
+
+        ds.search_async.assert_awaited_once_with(["q"], 3, {})
+        assert result == {"documents": expected}
+
 
 class TestChromaEmbeddingRetriever:
     def test_init(self, request):
@@ -165,3 +191,47 @@ def test_from_dict(self, request):
         assert retriever.filters == {"bar": "baz"}
         assert retriever.top_k == 42
         assert retriever.filter_policy == FilterPolicy.REPLACE
+
+    def test_from_dict_no_filter_policy(self, request):
+        data = {
+            "type": "haystack_integrations.components.retrievers.chroma.retriever.ChromaEmbeddingRetriever",
+            "init_parameters": {
+                "filters": {"bar": "baz"},
+                "top_k": 42,
+                "document_store": {
+                    "type": "haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore",
+                    "init_parameters": {
+                        "collection_name": "test_from_dict_no_filter_policy",
+                        "embedding_function": "HuggingFaceEmbeddingFunction",
+                        "persist_path": ".",
+                        "api_key": "1234567890",
+                        "distance_function": "l2",
+                    },
+                },
+            },
+        }
+        retriever = ChromaEmbeddingRetriever.from_dict(data)
+        assert retriever.filter_policy == FilterPolicy.REPLACE
+
+    def test_run_delegates_to_document_store_search_embeddings(self):
+        ds = mock.Mock(spec=ChromaDocumentStore)
+        expected = [Document(content="hit")]
+        ds.search_embeddings.return_value = [expected]
+        retriever = ChromaEmbeddingRetriever(ds, top_k=7)
+
+        result = retriever.run(query_embedding=[0.1, 0.2])
+
+        ds.search_embeddings.assert_called_once_with([[0.1, 0.2]], 7, {})
+        assert result == {"documents": expected}
+
+    @pytest.mark.asyncio
+    async def test_run_async_delegates_to_document_store_search_embeddings_async(self):
+        ds = mock.Mock(spec=ChromaDocumentStore)
+        expected = [Document(content="hit")]
+        ds.search_embeddings_async = mock.AsyncMock(return_value=[expected])
+        retriever = ChromaEmbeddingRetriever(ds, top_k=4)
+
+        result = await retriever.run_async(query_embedding=[0.5])
+
+        ds.search_embeddings_async.assert_awaited_once_with([[0.5]], 4, {})
+        assert result == {"documents": expected}