|
1 | 1 | import json |
2 | 2 | from pathlib import Path |
3 | | -from unittest.mock import patch |
| 3 | +from unittest.mock import Mock, patch |
4 | 4 |
|
5 | 5 | from openai.types.create_embedding_response import CreateEmbeddingResponse |
6 | 6 |
|
@@ -159,18 +159,14 @@ def test_fastembed_embeddings(): |
159 | 159 | assert_embedding_result(output) |
160 | 160 |
|
161 | 161 |
|
| 162 | +voyage_output_mock = Mock() |
| 163 | +voyage_output_mock.embeddings = [[1.0, 2.1, 3.2]] |
| 164 | + |
| 165 | + |
162 | 166 | @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") |
173 | 171 | 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