test: add unit tests for falkordb (#3283)

Author: anakin87

integrations/falkordb/tests/test_document_store.py

@@ -4,16 +4,25 @@
 
 import logging
 import os
+from types import SimpleNamespace
+from unittest.mock import MagicMock
 
+import falkordb as _falkordb_module
 import pytest
 from haystack.dataclasses import Document
+from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
+from haystack.document_stores.types import DuplicatePolicy
+from haystack.errors import FilterError
 from haystack.testing.document_store import DocumentStoreBaseTests
 
 from haystack_integrations.components.retrievers.falkordb import (
     FalkorDBCypherRetriever,
     FalkorDBEmbeddingRetriever,
 )
 from haystack_integrations.document_stores.falkordb import FalkorDBDocumentStore
+from haystack_integrations.document_stores.falkordb.document_store import (
+    _convert_filters,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +53,221 @@ def test_to_dict_from_dict(self):
         assert restored.similarity == "euclidean"
 
 
+def _result(rows):
+    return MagicMock(result_set=rows)
+
+
+@pytest.fixture
+def mock_falkordb(monkeypatch):
+    constructor = MagicMock()
+    client = MagicMock()
+    graph = MagicMock()
+    constructor.return_value = client
+    client.select_graph.return_value = graph
+    graph.query.return_value = _result([])
+    monkeypatch.setattr(_falkordb_module, "FalkorDB", constructor)
+    return constructor, client, graph
+
+
+class TestFalkorDBDocumentStoreUnit:
+    def test_init_rejects_invalid_similarity(self):
+        with pytest.raises(ValueError, match="not supported"):
+            FalkorDBDocumentStore(similarity="invalid")
+
+    @pytest.mark.parametrize(
+        "filter_node, expected_clause, expected_params",
+        [
+            ({"field": "year", "operator": "==", "value": 2024}, "coalesce(d.year = $p0, false)", {"p0": 2024}),
+            ({"field": "year", "operator": "==", "value": None}, "d.year IS NULL", {}),
+            ({"field": "y", "operator": "!=", "value": None}, "d.y IS NOT NULL", {}),
+            ({"field": "y", "operator": ">", "value": 1}, "coalesce(d.y > $p0, false)", {"p0": 1}),
+            ({"field": "y", "operator": ">", "value": None}, "false", {}),
+            (
+                {"field": "tag", "operator": "in", "value": ["a"]},
+                "coalesce(d.tag IN $p0, false)",
+                {"p0": ["a"]},
+            ),
+            (
+                {"field": "tag", "operator": "not in", "value": ["a"]},
+                "coalesce(NOT (d.tag IN $p0), true)",
+                {"p0": ["a"]},
+            ),
+            (
+                {"field": "meta.year", "operator": "==", "value": 2024},
+                "coalesce(d.year = $p0, false)",
+                {"p0": 2024},
+            ),
+            (
+                {
+                    "operator": "OR",
+                    "conditions": [
+                        {"field": "y", "operator": "==", "value": 1},
+                        {"operator": "NOT", "conditions": [{"field": "b", "operator": "<=", "value": 0}]},
+                    ],
+                },
+                "(coalesce(d.y = $p0, false) OR NOT (coalesce(d.b <= $p1, false)))",
+                {"p0": 1, "p1": 0},
+            ),
+        ],
+    )
+    def test_convert_filters(self, filter_node, expected_clause, expected_params):
+        clause, params = _convert_filters(filter_node)
+        assert clause == expected_clause
+        assert params == expected_params
+
+    @pytest.mark.parametrize(
+        "filter_node, match",
+        [
+            ({"operator": "AND"}, "requires a 'conditions' key"),
+            ({"operator": "NOT"}, "requires a 'conditions' key"),
+            ({"operator": "==", "value": 1}, "requires a 'field' key"),
+            ({"operator": "==", "field": "f"}, "requires a 'value' key"),
+            ({"field": "x", "operator": ">", "value": [1, 2]}, "does not support list values"),
+            ({"field": "x", "operator": ">", "value": "not-a-date"}, "non-ISO string"),
+            ({"field": "x", "operator": "in", "value": "scalar"}, "requires a list value"),
+            ({"field": "x", "operator": "not in", "value": "scalar"}, "requires a list value"),
+            ({"field": "x", "operator": "regex", "value": "."}, "Unsupported filter operator"),
+        ],
+    )
+    def test_convert_filters_errors(self, filter_node, match):
+        with pytest.raises(FilterError, match=match):
+            _convert_filters(filter_node)
+
+    @pytest.mark.parametrize("rows, expected", [([[42]], 42), ([], 0)])
+    def test_count_documents(self, mock_falkordb, rows, expected):
+        _, _, graph = mock_falkordb
+        graph.query.side_effect = [_result([]), _result([]), _result(rows)]
+        assert FalkorDBDocumentStore().count_documents() == expected
+
+    def test_filter_documents_no_filters(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        node = SimpleNamespace(properties={"id": "n1", "content": "hello"})
+        graph.query.side_effect = [_result([]), _result([]), _result([[node]])]
+        docs = FalkorDBDocumentStore().filter_documents()
+        assert [d.content for d in docs] == ["hello"]
+        assert "WHERE" not in graph.query.call_args_list[-1].args[0]
+
+    def test_filter_documents_with_filters(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        graph.query.side_effect = [_result([]), _result([]), _result([])]
+        FalkorDBDocumentStore().filter_documents({"field": "year", "operator": "==", "value": 2024})
+        last = graph.query.call_args_list[-1]
+        assert "WHERE" in last.args[0]
+        assert last.args[1] == {"p0": 2024}
+
+    @pytest.mark.usefixtures("mock_falkordb")
+    def test_filter_documents_malformed_raises(self):
+        with pytest.raises(FilterError, match="Invalid filter syntax"):
+            FalkorDBDocumentStore().filter_documents({"field": "year", "value": 2024})
+
+    def test_delete_documents_empty_is_noop(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        FalkorDBDocumentStore().delete_documents([])
+        assert graph.query.call_count == 2
+
+    def test_delete_documents_runs_query(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        FalkorDBDocumentStore().delete_documents(["a", "b"])
+        last = graph.query.call_args_list[-1]
+        assert "DETACH DELETE" in last.args[0]
+        assert last.args[1] == {"ids": ["a", "b"]}
+
+    @pytest.mark.parametrize(
+        "similarity, raw, scale_score, filters, expected_score, where_expected",
+        [
+            ("cosine", 0.4, True, None, 0.8, False),
+            ("cosine", 0.4, False, None, 0.4, False),
+            ("euclidean", 1.0, True, {"field": "y", "operator": "==", "value": 1}, 0.5, True),
+        ],
+    )
+    def test_embedding_retrieval(
+        self, mock_falkordb, similarity, raw, scale_score, filters, expected_score, where_expected
+    ):
+        _, _, graph = mock_falkordb
+        node = SimpleNamespace(properties={"id": "n1", "content": "hello"})
+        graph.query.side_effect = [_result([]), _result([]), _result([[node, raw]])]
+        docs = FalkorDBDocumentStore(similarity=similarity)._embedding_retrieval(
+            query_embedding=[0.1], top_k=5, filters=filters, scale_score=scale_score
+        )
+        assert docs[0].score == pytest.approx(expected_score)
+        assert ("WHERE" in graph.query.call_args_list[-1].args[0]) is where_expected
+
+    def test_cypher_retrieval_returns_documents(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        node = SimpleNamespace(properties={"id": "n1", "content": "hello"})
+        graph.query.side_effect = [_result([]), _result([]), _result([[node]])]
+        docs = FalkorDBDocumentStore()._cypher_retrieval("MATCH (d) RETURN d", parameters={"k": 1})
+        assert docs[0].content == "hello"
+
+    @pytest.mark.usefixtures("mock_falkordb")
+    def test_write_documents_rejects_non_documents(self):
+        with pytest.raises(ValueError, match="expects a list of Documents"):
+            FalkorDBDocumentStore().write_documents(["not a doc"])
+
+    @pytest.mark.usefixtures("mock_falkordb")
+    def test_write_documents_empty_list_returns_zero(self, caplog):
+        with caplog.at_level(logging.WARNING):
+            assert FalkorDBDocumentStore().write_documents([]) == 0
+        assert "empty list" in caplog.text
+
+    def test_write_documents_policy_none_coerced_to_fail(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        graph.query.side_effect = [_result([]), _result([]), _result([["a"]])]
+        with pytest.raises(DuplicateDocumentError):
+            FalkorDBDocumentStore().write_documents([Document(id="a", content="x")])
+
+    def test_write_documents_drops_duplicates_within_batch(self, mock_falkordb, caplog):
+        _, _, graph = mock_falkordb
+        graph.query.side_effect = [_result([]), _result([]), _result([]), _result([[2]])]
+        with caplog.at_level(logging.INFO):
+            written = FalkorDBDocumentStore().write_documents(
+                [
+                    Document(id="a", content="x"),
+                    Document(id="a", content="x"),
+                    Document(id="b", content="y"),
+                ],
+                policy=DuplicatePolicy.SKIP,
+            )
+        assert written == 2
+        assert "already present in the batch" in caplog.text
+        sent_ids = [d["id"] for d in graph.query.call_args_list[-1].args[1]["docs"]]
+        assert sent_ids == ["a", "b"]
+
+    def test_write_documents_skip_filters_existing(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        graph.query.side_effect = [_result([]), _result([]), _result([["a"]]), _result([[1]])]
+        written = FalkorDBDocumentStore().write_documents(
+            [Document(id="a", content="x"), Document(id="b", content="y")],
+            policy=DuplicatePolicy.SKIP,
+        )
+        assert written == 1
+        sent_ids = [d["id"] for d in graph.query.call_args_list[-1].args[1]["docs"]]
+        assert sent_ids == ["b"]
+
+    def test_write_documents_overwrite_uses_on_match_set(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        graph.query.side_effect = [_result([]), _result([]), _result([[1]])]
+        FalkorDBDocumentStore().write_documents([Document(id="a", content="x")], policy=DuplicatePolicy.OVERWRITE)
+        assert "ON MATCH SET d = doc" in graph.query.call_args_list[-1].args[0]
+
+    def test_write_documents_embeddings_second_pass(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        graph.query.side_effect = [_result([]), _result([]), _result([[1]]), _result([])]
+        FalkorDBDocumentStore().write_documents(
+            [Document(id="a", content="x", embedding=[0.1, 0.2, 0.3])],
+            policy=DuplicatePolicy.OVERWRITE,
+        )
+        last = graph.query.call_args_list[-1]
+        assert "vecf32" in last.args[0]
+        assert last.args[1]["docs"] == [{"id": "a", "emb": [0.1, 0.2, 0.3]}]
+
+    def test_write_documents_wraps_errors(self, mock_falkordb):
+        _, _, graph = mock_falkordb
+        graph.query.side_effect = [_result([]), _result([]), Exception("boom")]
+        with pytest.raises(DocumentStoreError, match="Failed to write documents"):
+            FalkorDBDocumentStore().write_documents([Document(id="a", content="x")], policy=DuplicatePolicy.OVERWRITE)
+
+
 @pytest.mark.integration
 class TestDocumentStore(DocumentStoreBaseTests):
     """