From 892d1a378d772faaad2e69267dbe220cf4fabe99 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Mon, 16 Jun 2025 18:10:03 +0200
Subject: [PATCH] test search

---
 tests/test_chunk.py  | 29 -------------------------
 tests/test_search.py | 50 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 29 deletions(-)
 create mode 100644 tests/test_search.py

diff --git a/tests/test_chunk.py b/tests/test_chunk.py
index a9263c7c..0fbd3eac 100644
--- a/tests/test_chunk.py
+++ b/tests/test_chunk.py
@@ -8,35 +8,6 @@ from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
 
 
-@pytest.mark.asyncio
-async def test_search_chunks(qa_corpus: Dataset):
-    """Test vector search functionality using ChunkRepository."""
-    # Create an in-memory store and repositories
-    store = Store(":memory:")
-    doc_repo = DocumentRepository(store)
-    chunk_repo = ChunkRepository(store)
-
-    # Get the first document from the corpus
-    first_doc = qa_corpus[0]
-    document_text = first_doc["document_extracted"]
-
-    # Create and store a document
-    document = Document(content=document_text, metadata={"source": "qa_corpus"})
-    created_document = await doc_repo.create(document)
-
-    # Perform a search using ChunkRepository
-    search_query = "news"  # Simple query
-    results = await chunk_repo.search_chunks(search_query, limit=3)
-
-    # Verify search results
-    assert len(results) <= 3
-    assert all(hasattr(chunk, "content") for chunk in results)
-    assert all(hasattr(chunk, "document_id") for chunk in results)
-    assert all(chunk.document_id == created_document.id for chunk in results)
-
-    store.close()
-
-
 @pytest.mark.asyncio
 async def test_chunk_repository_operations(qa_corpus: Dataset):
     """Test ChunkRepository operations."""
diff --git a/tests/test_search.py b/tests/test_search.py
new file mode 100644
index 00000000..1f23717e
--- /dev/null
+++ b/tests/test_search.py
@@ -0,0 +1,50 @@
+import pytest
+from datasets import Dataset
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.models.document import Document
+from haiku.rag.store.repositories.chunk import ChunkRepository
+from haiku.rag.store.repositories.document import DocumentRepository
+
+
+@pytest.mark.asyncio
+async def test_search_qa_corpus(qa_corpus: Dataset):
+    """Test that documents can be found by searching with their associated questions."""
+    # Create an in-memory store and repositories
+    store = Store(":memory:")
+    doc_repo = DocumentRepository(store)
+    chunk_repo = ChunkRepository(store)
+    num_documents = 10
+    # Load first 10 documents with embeddings (reduced for faster testing)
+    documents = []
+    for i in range(num_documents):
+        doc_data = qa_corpus[i]
+        document_text = doc_data["document_extracted"]
+
+        # Create a Document instance
+        document = Document(
+            content=document_text,
+            metadata={
+                "source": "qa_corpus",
+                "topic": doc_data.get("document_topic", ""),
+                "document_id": doc_data.get("document_id", ""),
+                "question": doc_data["question"],
+            },
+        )
+
+        # Create the document with chunks and embeddings
+        created_document = await doc_repo.create(document)
+        documents.append((created_document, doc_data))
+
+    for i in range(num_documents):  # Test with first few documents
+        target_document, doc_data = documents[i]
+        question = doc_data["question"]
+
+        # Search for chunks using the question
+        search_results = await chunk_repo.search_chunks(question, limit=5)
+
+        # Check if target document is in results
+        target_document_ids = {chunk.document_id for chunk in search_results}
+        assert target_document.id in target_document_ids
+
+    store.close()