From b0d806526a7f5218d1cafa485c8dc502552026d6 Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 10 Dec 2025 11:47:29 +0100
Subject: [PATCH] feat: adds missing checks to rag completion test

---
 .../retrieval/test_rag_completion_retriever.py     | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/cognee/tests/integration/retrieval/test_rag_completion_retriever.py b/cognee/tests/integration/retrieval/test_rag_completion_retriever.py
index 5ee3812a0..b01d58160 100644
--- a/cognee/tests/integration/retrieval/test_rag_completion_retriever.py
+++ b/cognee/tests/integration/retrieval/test_rag_completion_retriever.py
@@ -211,7 +211,19 @@ async def test_rag_completion_context_simple(setup_test_environment_with_chunks_
 
     context = await retriever.get_context("Mike")
 
-    assert context == "Mike Broski", "Failed to get Mike Broski"
+    assert isinstance(context, str), "Context should be a string"
+    assert "Mike Broski" in context, "Failed to get Mike Broski"
+
+
+@pytest.mark.asyncio
+async def test_rag_completion_context_multiple_chunks(setup_test_environment_with_chunks_simple):
+    """Integration test: verify CompletionRetriever can retrieve context from multiple chunks."""
+    retriever = CompletionRetriever()
+
+    context = await retriever.get_context("Steve")
+
+    assert isinstance(context, str), "Context should be a string"
+    assert "Steve Rodger" in context, "Failed to get Steve Rodger"
 
 
 @pytest.mark.asyncio