Fix tests

deepsense-ai · Oct 7, 2024 · 227fddc · 227fddc
1 parent d306f84
commit 227fddc
Showing 1 changed file with 8 additions and 20 deletions.
diff --git a/packages/ragbits-document-search/tests/unit/test_document_search.py b/packages/ragbits-document-search/tests/unit/test_document_search.py
@@ -4,7 +4,6 @@
 
 import pytest
 
-from ragbits.core.embeddings.noop import NoopEmbeddings
 from ragbits.core.vector_store.in_memory import InMemoryVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search._main import SearchConfig
@@ -19,27 +18,16 @@
 }
 
 
-async def test_document_search():
-    document_search = DocumentSearch(embedder=NoopEmbeddings(), vector_store=InMemoryVectorStore())
-
-    await document_search.ingest_document(
-        DocumentMeta.create_text_document_from_literal("Name of Peppa's brother is George")
-    )
-
-    results = await document_search.search("Peppa's brother")
-
-    first_result = results[0]
-
-    assert isinstance(first_result, TextElement)
-    assert first_result.content == "Name of Peppa's brother is George"
-
-
-async def test_document_search_from_config():
+@pytest.mark.parametrize(
+    "document",
+    [
+        DocumentMeta.create_text_document_from_literal("Name of Peppa's brother is George"),
+    ],
+)
+async def test_document_search_from_config(document):
     document_search = DocumentSearch.from_config(CONFIG)
 
-    await document_search.ingest_document(
-        DocumentMeta.create_text_document_from_literal("Name of Peppa's brother is George")
-    )
+    await document_search.ingest_document(document, document_processor=DummyProvider())
     results = await document_search.search("Peppa's brother")
 
     first_result = results[0]