feat: adds unit and integration test to triplet retriever

cognee/tests/integration/retrieval/test_triplet_retriever.py

@@ -0,0 +1,84 @@
+import os
+import pytest
+import pathlib
+import pytest_asyncio
+import cognee
+
+from cognee.low_level import setup
+from cognee.tasks.storage import add_data_points
+from cognee.modules.retrieval.exceptions.exceptions import NoDataError
+from cognee.modules.retrieval.triplet_retriever import TripletRetriever
+from cognee.modules.engine.models import Triplet
+
+
+@pytest_asyncio.fixture
+async def setup_test_environment_with_triplets():
+    """Set up a clean test environment with triplets."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(base_dir / ".cognee_system/test_triplet_retriever_context_simple")
+    data_directory_path = str(base_dir / ".data_storage/test_triplet_retriever_context_simple")
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+    await setup()
+
+    triplet1 = Triplet(
+        from_node_id="node1",
+        to_node_id="node2",
+        text="Alice knows Bob",
+    )
+    triplet2 = Triplet(
+        from_node_id="node2",
+        to_node_id="node3",
+        text="Bob works at Tech Corp",
+    )
+
+    triplets = [triplet1, triplet2]
+    await add_data_points(triplets)
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest_asyncio.fixture
+async def setup_test_environment_empty():
+    """Set up a clean test environment without triplets."""
+    base_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    system_directory_path = str(
+        base_dir / ".cognee_system/test_triplet_retriever_context_empty_collection"
+    )
+    data_directory_path = str(
+        base_dir / ".data_storage/test_triplet_retriever_context_empty_collection"
+    )
+
+    cognee.config.system_root_directory(system_directory_path)
+    cognee.config.data_root_directory(data_directory_path)
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    yield
+
+    try:
+        await cognee.prune.prune_data()
+        await cognee.prune.prune_system(metadata=True)
+    except Exception:
+        pass
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_context_simple(setup_test_environment_with_triplets):
+    """Integration test: verify TripletRetriever can retrieve triplet context."""
+    retriever = TripletRetriever(top_k=5)
+
+    context = await retriever.get_context("Alice")
+
+    assert "Alice knows Bob" in context, "Failed to get Alice triplet"

cognee/tests/unit/modules/retrieval/triplet_retriever_test.py

@@ -0,0 +1,83 @@
+import pytest
+from unittest.mock import AsyncMock, patch, MagicMock
+
+from cognee.modules.retrieval.triplet_retriever import TripletRetriever
+from cognee.modules.retrieval.exceptions.exceptions import NoDataError
+from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
+
+
+@pytest.fixture
+def mock_vector_engine():
+    """Create a mock vector engine."""
+    engine = AsyncMock()
+    engine.has_collection = AsyncMock(return_value=True)
+    engine.search = AsyncMock()
+    return engine
+
+
+@pytest.mark.asyncio
+async def test_get_context_success(mock_vector_engine):
+    """Test successful retrieval of triplet context."""
+    mock_result1 = MagicMock()
+    mock_result1.payload = {"text": "Alice knows Bob"}
+    mock_result2 = MagicMock()
+    mock_result2.payload = {"text": "Bob works at Tech Corp"}
+
+    mock_vector_engine.search.return_value = [mock_result1, mock_result2]
+
+    retriever = TripletRetriever(top_k=5)
+
+    with patch(
+        "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        context = await retriever.get_context("test query")
+
+    assert context == "Alice knows Bob\nBob works at Tech Corp"
+    mock_vector_engine.search.assert_awaited_once_with("Triplet_text", "test query", limit=5)
+
+
+@pytest.mark.asyncio
+async def test_get_context_no_collection(mock_vector_engine):
+    """Test that NoDataError is raised when Triplet_text collection doesn't exist."""
+    mock_vector_engine.has_collection.return_value = False
+
+    retriever = TripletRetriever()
+
+    with patch(
+        "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        with pytest.raises(NoDataError, match="create_triplet_embeddings"):
+            await retriever.get_context("test query")
+
+
+@pytest.mark.asyncio
+async def test_get_context_empty_results(mock_vector_engine):
+    """Test that empty string is returned when no triplets are found."""
+    mock_vector_engine.search.return_value = []
+
+    retriever = TripletRetriever()
+
+    with patch(
+        "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        context = await retriever.get_context("test query")
+
+    assert context == ""
+
+
+@pytest.mark.asyncio
+async def test_get_context_collection_not_found_error(mock_vector_engine):
+    """Test that CollectionNotFoundError is converted to NoDataError."""
+    mock_vector_engine.has_collection.side_effect = CollectionNotFoundError("Collection not found")
+
+    retriever = TripletRetriever()
+
+    with patch(
+        "cognee.modules.retrieval.triplet_retriever.get_vector_engine",
+        return_value=mock_vector_engine,
+    ):
+        with pytest.raises(NoDataError, match="No data found"):
+            await retriever.get_context("test query")