Fix tests

awinml · Dec 11, 2023 · 61f852c · 61f852c
1 parent e2bedf7
commit 61f852c
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/src/voyage_embedders/voyage_document_embedder.py b/src/voyage_embedders/voyage_document_embedder.py
@@ -2,9 +2,9 @@
 from typing import Any, Dict, List, Optional
 
 import voyageai
-from haystack.dataclasses import Document
 from haystack.core.component import component
 from haystack.core.serialization import default_to_dict
+from haystack.dataclasses import Document
 from tqdm import tqdm
 from voyageai import get_embeddings
 

diff --git a/tests/test_voyage_document_embedder.py b/tests/test_voyage_document_embedder.py
@@ -65,7 +65,7 @@ def test_to_dict(self):
         component = VoyageDocumentEmbedder(api_key="fake-api-key")
         data = component.to_dict()
         assert data == {
-            "type": "VoyageDocumentEmbedder",
+            "type": "voyage_embedders.voyage_document_embedder.VoyageDocumentEmbedder",
             "init_parameters": {
                 "model_name": "voyage-01",
                 "prefix": "",
@@ -91,7 +91,7 @@ def test_to_dict_with_custom_init_parameters(self):
         )
         data = component.to_dict()
         assert data == {
-            "type": "VoyageDocumentEmbedder",
+            "type": "voyage_embedders.voyage_document_embedder.VoyageDocumentEmbedder",
             "init_parameters": {
                 "model_name": "model",
                 "prefix": "prefix",
@@ -187,6 +187,7 @@ def test_run(self):
                     "prefix ML | A transformer is a deep learning architecture suffix",
                 ],
                 batch_size=8,
+                input_type="document",
             )
         documents_with_embeddings = result["documents"]
 

diff --git a/tests/test_voyage_text_embedder.py b/tests/test_voyage_text_embedder.py
@@ -53,7 +53,7 @@ def test_to_dict(self):
         component = VoyageTextEmbedder(api_key="fake-api-key")
         data = component.to_dict()
         assert data == {
-            "type": "VoyageTextEmbedder",
+            "type": "voyage_embedders.voyage_text_embedder.VoyageTextEmbedder",
             "init_parameters": {
                 "model_name": "voyage-01",
                 "input_type": "query",
@@ -73,7 +73,7 @@ def test_to_dict_with_custom_init_parameters(self):
         )
         data = component.to_dict()
         assert data == {
-            "type": "VoyageTextEmbedder",
+            "type": "voyage_embedders.voyage_text_embedder.VoyageTextEmbedder",
             "init_parameters": {
                 "model_name": "model",
                 "input_type": "document",
@@ -92,7 +92,9 @@ def test_run(self):
             embedder = VoyageTextEmbedder(api_key="fake-api-key", model_name=model, prefix="prefix ", suffix=" suffix")
             result = embedder.run(text="The food was delicious")
 
-            voyageai_embedding_patch.assert_called_once_with(model=model, text="prefix The food was delicious suffix")
+            voyageai_embedding_patch.assert_called_once_with(
+                model=model, text="prefix The food was delicious suffix", input_type="query"
+            )
 
         assert len(result["embedding"]) == 1024
         assert all(isinstance(x, float) for x in result["embedding"])