Skip to content

Commit 07408e9

Browse files
committed
fix: update test case
1 parent 0146b8b commit 07408e9

File tree

1 file changed

+10
-14
lines changed

1 file changed

+10
-14
lines changed

libs/kotaemon/tests/test_embedding_models.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import json
22
from pathlib import Path
3-
from unittest.mock import patch
3+
from unittest.mock import Mock, patch
44

55
from openai.types.create_embedding_response import CreateEmbeddingResponse
66

@@ -159,18 +159,14 @@ def test_fastembed_embeddings():
159159
assert_embedding_result(output)
160160

161161

162+
voyage_output_mock = Mock()
163+
voyage_output_mock.embeddings = [[1.0, 2.1, 3.2]]
164+
165+
162166
@skip_when_voyageai_not_installed
163-
@patch(
164-
"voyageai.Client.embed",
165-
side_effect=lambda *args, **kwargs: [[1.0, 2.1, 3.2]],
166-
)
167-
@patch(
168-
"voyageai.AsyncClient.embed",
169-
side_effect=lambda *args, **kwargs: [[1.0, 2.1, 3.2]],
170-
)
171-
def test_voyageai_embeddings():
172-
model = VoyageAIEmbeddings()
167+
@patch("voyageai.Client.embed", return_value=voyage_output_mock)
168+
@patch("voyageai.AsyncClient.embed", return_value=voyage_output_mock)
169+
def test_voyageai_embeddings(sync_call, async_call):
170+
model = VoyageAIEmbeddings(api_key="test")
173171
output = model("Hello, world!")
174-
assert isinstance(output, list) and all(
175-
isinstance(doc, DocumentWithEmbedding) for doc in output
176-
)
172+
assert all(isinstance(doc, DocumentWithEmbedding) for doc in output)

0 commit comments

Comments
 (0)