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")