fix: mock embedder in integration tests

deepsense-ai · Nov 28, 2024 · c8870ab · c8870ab
1 parent af80819
commit c8870ab
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 12 deletions.
diff --git a/packages/ragbits-core/tests/integration/vector_stores/test_chroma.py b/packages/ragbits-core/tests/integration/vector_stores/test_chroma.py
@@ -1,10 +1,10 @@
 import os
 import tempfile
 from pathlib import Path
+from unittest.mock import AsyncMock
 
 from chromadb import EphemeralClient
 
-from ragbits.core.embeddings.litellm import LiteLLMEmbeddings
 from ragbits.core.vector_stores.chroma import ChromaVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
@@ -26,9 +26,8 @@ async def test_update_document() -> None:
     document_1 = DocumentMeta.from_local_path(Path(document_1_path))
     document_2 = DocumentMeta.from_local_path(Path(document_2_path))
 
-    embedder = LiteLLMEmbeddings(
-        model="text-embedding-3-small",
-    )
+    embedder = AsyncMock()
+    embedder.embed_text.return_value = [[0.0], [0.0]]
     vector_store = ChromaVectorStore(
         client=EphemeralClient(),
         index_name="test_index_name",

diff --git a/packages/ragbits-core/tests/integration/vector_stores/test_in_memory.py b/packages/ragbits-core/tests/integration/vector_stores/test_in_memory.py
@@ -1,8 +1,8 @@
 import os
 import tempfile
 from pathlib import Path
+from unittest.mock import AsyncMock
 
-from ragbits.core.embeddings.litellm import LiteLLMEmbeddings
 from ragbits.core.vector_stores.in_memory import InMemoryVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
@@ -24,9 +24,8 @@ async def test_update_document() -> None:
     document_1 = DocumentMeta.from_local_path(Path(document_1_path))
     document_2 = DocumentMeta.from_local_path(Path(document_2_path))
 
-    embedder = LiteLLMEmbeddings(
-        model="text-embedding-3-small",
-    )
+    embedder = AsyncMock()
+    embedder.embed_text.return_value = [[0.0], [0.0]]
     vector_store = InMemoryVectorStore()
     document_search = DocumentSearch(
         embedder=embedder,

diff --git a/packages/ragbits-core/tests/integration/vector_stores/test_qdrant.py b/packages/ragbits-core/tests/integration/vector_stores/test_qdrant.py
@@ -1,10 +1,10 @@
 import os
 import tempfile
 from pathlib import Path
+from unittest.mock import AsyncMock
 
 from qdrant_client import AsyncQdrantClient
 
-from ragbits.core.embeddings.litellm import LiteLLMEmbeddings
 from ragbits.core.vector_stores.qdrant import QdrantVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
@@ -26,9 +26,8 @@ async def test_update_document() -> None:
     document_1 = DocumentMeta.from_local_path(Path(document_1_path))
     document_2 = DocumentMeta.from_local_path(Path(document_2_path))
 
-    embedder = LiteLLMEmbeddings(
-        model="text-embedding-3-small",
-    )
+    embedder = AsyncMock()
+    embedder.embed_text.return_value = [[0.0], [0.0]]
     vector_store = QdrantVectorStore(
         client=AsyncQdrantClient(":memory:"),
         index_name="test_index_name",