From 714eeaffc432d1ade597af7f108a1a09f4fc1794 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 10 Dec 2025 11:48:18 +0100
Subject: [PATCH] feat: adds some new tests to the triplet retriever int test

---
 .../retrieval/test_triplet_retriever.py       | 35 +++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/cognee/tests/integration/retrieval/test_triplet_retriever.py b/cognee/tests/integration/retrieval/test_triplet_retriever.py
index e547b6cbe..ebe853e08 100644
--- a/cognee/tests/integration/retrieval/test_triplet_retriever.py
+++ b/cognee/tests/integration/retrieval/test_triplet_retriever.py
@@ -82,3 +82,38 @@ async def test_triplet_retriever_context_simple(setup_test_environment_with_trip
     context = await retriever.get_context("Alice")
 
     assert "Alice knows Bob" in context, "Failed to get Alice triplet"
+    assert isinstance(context, str), "Context should be a string"
+    assert len(context) > 0, "Context should not be empty"
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_context_multiple_triplets(setup_test_environment_with_triplets):
+    """Integration test: verify TripletRetriever can retrieve multiple triplets."""
+    retriever = TripletRetriever(top_k=5)
+
+    context = await retriever.get_context("Bob")
+
+    assert "Alice knows Bob" in context or "Bob works at Tech Corp" in context, (
+        "Failed to get Bob-related triplets"
+    )
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_top_k_limit(setup_test_environment_with_triplets):
+    """Integration test: verify TripletRetriever respects top_k parameter."""
+    retriever = TripletRetriever(top_k=1)
+
+    context = await retriever.get_context("Alice")
+
+    assert isinstance(context, str), "Context should be a string"
+
+
+@pytest.mark.asyncio
+async def test_triplet_retriever_context_empty(setup_test_environment_empty):
+    """Integration test: verify TripletRetriever handles empty graph correctly."""
+    await setup()
+
+    retriever = TripletRetriever()
+
+    with pytest.raises(NoDataError):
+        await retriever.get_context("Alice")