Skip to content

Commit b6299b9

Browse files
omborda2002bogdankostic
omborda2002
and
bogdankostic
authored
test: type-check test/components/embedders and add it to mypy config (#11859)
Co-authored-by: bogdankostic <bogdankostic@web.de>
1 parent 11ec589 commit b6299b9

6 files changed

Lines changed: 44 additions & 16 deletions

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ all = 'pytest {args:test}'
149149
e2e = 'pytest {args:e2e}'
150150

151151
# TODO We want to eventually type the whole test folder
152-
types = "mypy --install-types --non-interactive --cache-dir=.mypy_cache/ {args:haystack test/core/ test/marshal/ test/testing/ test/tracing/ test/tools/ test/human_in_the_loop test/evaluation test/document_stores test/dataclasses test/utils/}"
152+
types = "mypy --install-types --non-interactive --cache-dir=.mypy_cache/ {args:haystack test/core/ test/marshal/ test/testing/ test/tracing/ test/tools/ test/human_in_the_loop test/evaluation test/document_stores test/dataclasses test/utils/ test/components/embedders/}"
153153

154154
[project.urls]
155155
"CI: GitHub" = "https://github.com/deepset-ai/haystack/actions"

‎test/components/embedders/test_azure_document_embedder.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ def test_embed_batch_handles_exceptions_gracefully(self, caplog):
216216
embedding_separator=" | ",
217217
)
218218
embedder.warm_up()
219+
assert embedder.client is not None
219220

220221
fake_texts_to_embed = {"1": "text1", "2": "text2"}
221222

@@ -237,6 +238,7 @@ def test_embed_batch_raises_exception_on_failure(self):
237238
raise_on_failure=True,
238239
)
239240
embedder.warm_up()
241+
assert embedder.client is not None
240242
fake_texts_to_embed = {"1": "text1", "2": "text2"}
241243
with patch.object(
242244
embedder.client.embeddings,
@@ -303,6 +305,7 @@ def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
303305
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
304306
embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
305307
embedder.warm_up()
308+
assert embedder.client is not None
306309
assert embedder.client.max_retries == 5
307310
assert embedder.client.timeout == 30.0
308311

@@ -312,6 +315,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_parameters(self, monkeypatch)
312315
azure_endpoint="https://example-resource.azure.openai.com/", timeout=40.0, max_retries=1
313316
)
314317
embedder.warm_up()
318+
assert embedder.client is not None
315319
assert embedder.client.max_retries == 1
316320
assert embedder.client.timeout == 40.0
317321

@@ -321,6 +325,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
321325
monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
322326
embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
323327
embedder.warm_up()
328+
assert embedder.client is not None
324329
assert embedder.client.max_retries == 10
325330
assert embedder.client.timeout == 100.0
326331

@@ -334,6 +339,7 @@ def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
334339

335340
def test_sync_lifecycle(self, mock_azure_clients):
336341
sync_cls, _ = mock_azure_clients
342+
sync_client = sync_cls.return_value
337343
embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
338344
assert embedder.client is None
339345
assert embedder.async_client is None
@@ -343,19 +349,20 @@ def test_sync_lifecycle(self, mock_azure_clients):
343349
assert embedder.async_client is None
344350

345351
embedder.close()
346-
sync_cls.return_value.close.assert_called_once()
352+
sync_client.close.assert_called_once()
347353
assert embedder.client is None
348354

349355
async def test_async_lifecycle(self, mock_azure_clients):
350356
_, async_cls = mock_azure_clients
357+
async_client = async_cls.return_value
351358
embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
352359

353360
await embedder.warm_up_async()
354361
assert embedder.async_client is async_cls.return_value
355362
assert embedder.client is None
356363

357364
await embedder.close_async()
358-
async_cls.return_value.close.assert_awaited_once()
365+
async_client.close.assert_awaited_once()
359366
assert embedder.async_client is None
360367

361368
async def test_close_is_safe_without_warm_up(self, mock_azure_clients):

‎test/components/embedders/test_azure_text_embedder.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
216216
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
217217
embedder = AzureOpenAITextEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
218218
embedder.warm_up()
219+
assert embedder.client is not None
219220
assert embedder.client.max_retries == 5
220221
assert embedder.client.timeout == 30.0
221222

@@ -225,6 +226,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_parameters(self, monkeypatch)
225226
azure_endpoint="https://example-resource.azure.openai.com/", timeout=40.0, max_retries=1
226227
)
227228
embedder.warm_up()
229+
assert embedder.client is not None
228230
assert embedder.client.max_retries == 1
229231
assert embedder.client.timeout == 40.0
230232

@@ -234,6 +236,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
234236
monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
235237
embedder = AzureOpenAITextEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
236238
embedder.warm_up()
239+
assert embedder.client is not None
237240
assert embedder.client.max_retries == 10
238241
assert embedder.client.timeout == 100.0
239242

@@ -247,6 +250,7 @@ def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
247250

248251
def test_sync_lifecycle(self, mock_azure_clients):
249252
sync_cls, _ = mock_azure_clients
253+
sync_client = sync_cls.return_value
250254
embedder = AzureOpenAITextEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
251255
assert embedder.client is None
252256
assert embedder.async_client is None
@@ -256,19 +260,20 @@ def test_sync_lifecycle(self, mock_azure_clients):
256260
assert embedder.async_client is None
257261

258262
embedder.close()
259-
sync_cls.return_value.close.assert_called_once()
263+
sync_client.close.assert_called_once()
260264
assert embedder.client is None
261265

262266
async def test_async_lifecycle(self, mock_azure_clients):
263267
_, async_cls = mock_azure_clients
268+
async_client = async_cls.return_value
264269
embedder = AzureOpenAITextEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
265270

266271
await embedder.warm_up_async()
267272
assert embedder.async_client is async_cls.return_value
268273
assert embedder.client is None
269274

270275
await embedder.close_async()
271-
async_cls.return_value.close.assert_awaited_once()
276+
async_client.close.assert_awaited_once()
272277
assert embedder.async_client is None
273278

274279
async def test_close_is_safe_without_warm_up(self, mock_azure_clients):

‎test/components/embedders/test_mock_text_embedder.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,8 @@ def test_embedding_fn(self):
6060
assert MockTextEmbedder(embedding_fn=_ones).run("hello")["embedding"] == [1.0, 1.0, 1.0]
6161

6262
def test_embedding_fn_invalid_return_raises(self):
63-
embedder = MockTextEmbedder(embedding_fn=lambda text: "not a vector")
63+
# embedding_fn deliberately returns a non-vector to exercise the runtime type check
64+
embedder = MockTextEmbedder(embedding_fn=lambda text: "not a vector") # type: ignore[arg-type, return-value]
6465
with pytest.raises(TypeError, match="must be a sequence of numbers"):
6566
embedder.run("hello")
6667

@@ -79,8 +80,9 @@ def test_meta(self):
7980
assert custom["extra"] == "value"
8081

8182
def test_run_rejects_non_string(self):
83+
# a non-string input is passed on purpose to exercise the runtime type check
8284
with pytest.raises(TypeError, match="expects a string"):
83-
MockTextEmbedder().run(["not", "a", "string"])
85+
MockTextEmbedder().run(["not", "a", "string"]) # type: ignore[arg-type]
8486

8587
async def test_run_async(self):
8688
embedder = MockTextEmbedder(dimension=8)

‎test/components/embedders/test_openai_document_embedder.py‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -200,15 +200,15 @@ def test_run_wrong_input_format(self):
200200
list_integers_input = [1, 2, 3]
201201

202202
with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"):
203-
embedder.run(documents=string_input)
203+
embedder.run(documents=string_input) # type: ignore[arg-type]
204204

205205
with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"):
206-
embedder.run(documents=list_integers_input)
206+
embedder.run(documents=list_integers_input) # type: ignore[arg-type]
207207

208208
def test_run_on_empty_list(self):
209209
embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake-api-key"))
210210

211-
empty_list_input = []
211+
empty_list_input: list[Document] = []
212212
result = embedder.run(documents=empty_list_input)
213213

214214
assert result["documents"] is not None
@@ -217,6 +217,7 @@ def test_run_on_empty_list(self):
217217
def test_embed_batch_handles_exceptions_gracefully(self, caplog):
218218
embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake_api_key"))
219219
embedder.warm_up()
220+
assert embedder.client is not None
220221
fake_texts_to_embed = {"1": "text1", "2": "text2"}
221222
with patch.object(
222223
embedder.client.embeddings,
@@ -231,6 +232,7 @@ def test_embed_batch_handles_exceptions_gracefully(self, caplog):
231232
def test_run_handles_exceptions_gracefully(self, caplog):
232233
embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake_api_key"), batch_size=1)
233234
embedder.warm_up()
235+
assert embedder.client is not None
234236
docs = [
235237
Document(content="I love cheese", meta={"topic": "Cuisine"}),
236238
Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}),
@@ -260,6 +262,7 @@ def test_run_handles_exceptions_gracefully(self, caplog):
260262
def test_embed_batch_raises_exception_on_failure(self):
261263
embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake_api_key"), raise_on_failure=True)
262264
embedder.warm_up()
265+
assert embedder.client is not None
263266
fake_texts_to_embed = {"1": "text1", "2": "text2"}
264267
with patch.object(
265268
embedder.client.embeddings,
@@ -354,12 +357,14 @@ def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
354357
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
355358
embedder = OpenAIDocumentEmbedder()
356359
embedder.warm_up()
360+
assert embedder.client is not None
357361
assert embedder.client.max_retries == 5
358362
assert embedder.client.timeout == 30.0
359363

360364
def test_warm_up_uses_timeout_and_max_retries_from_parameters(self):
361365
embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake-api-key"), timeout=40.0, max_retries=1)
362366
embedder.warm_up()
367+
assert embedder.client is not None
363368
assert embedder.client.max_retries == 1
364369
assert embedder.client.timeout == 40.0
365370

@@ -368,6 +373,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
368373
monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
369374
embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake-api-key"))
370375
embedder.warm_up()
376+
assert embedder.client is not None
371377
assert embedder.client.max_retries == 10
372378
assert embedder.client.timeout == 100.0
373379

@@ -379,6 +385,7 @@ def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
379385

380386
def test_sync_lifecycle(self, mock_openai_clients):
381387
sync_cls, _ = mock_openai_clients
388+
sync_client = sync_cls.return_value
382389
embedder = OpenAIDocumentEmbedder()
383390
assert embedder.client is None
384391
assert embedder.async_client is None
@@ -388,19 +395,20 @@ def test_sync_lifecycle(self, mock_openai_clients):
388395
assert embedder.async_client is None
389396

390397
embedder.close()
391-
sync_cls.return_value.close.assert_called_once()
398+
sync_client.close.assert_called_once()
392399
assert embedder.client is None
393400

394401
async def test_async_lifecycle(self, mock_openai_clients):
395402
_, async_cls = mock_openai_clients
403+
async_client = async_cls.return_value
396404
embedder = OpenAIDocumentEmbedder()
397405

398406
await embedder.warm_up_async()
399407
assert embedder.async_client is async_cls.return_value
400408
assert embedder.client is None
401409

402410
await embedder.close_async()
403-
async_cls.return_value.close.assert_awaited_once()
411+
async_client.close.assert_awaited_once()
404412
assert embedder.async_client is None
405413

406414
async def test_close_is_safe_without_warm_up(self, mock_openai_clients):

‎test/components/embedders/test_openai_text_embedder.py‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
import pytest
1010
from openai.types import CreateEmbeddingResponse, Embedding
11+
from openai.types.create_embedding_response import Usage
1112

1213
import haystack.components.embedders.openai_text_embedder as openai_text_embedder_module
1314
from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder
@@ -169,7 +170,7 @@ def test_prepare_output(self, monkeypatch):
169170
data=[Embedding(embedding=[0.1, 0.2, 0.3], index=0, object="embedding")],
170171
model="text-embedding-ada-002",
171172
object="list",
172-
usage={"prompt_tokens": 6, "total_tokens": 6},
173+
usage=Usage(prompt_tokens=6, total_tokens=6),
173174
)
174175

175176
embedder = OpenAITextEmbedder()
@@ -185,7 +186,7 @@ def test_run_wrong_input_format(self):
185186
list_integers_input = [1, 2, 3]
186187

187188
with pytest.raises(TypeError, match="OpenAITextEmbedder expects a string as an input"):
188-
embedder.run(text=list_integers_input)
189+
embedder.run(text=list_integers_input) # type: ignore[arg-type]
189190

190191
@pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
191192
@pytest.mark.integration
@@ -241,12 +242,14 @@ def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
241242
monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
242243
embedder = OpenAITextEmbedder()
243244
embedder.warm_up()
245+
assert embedder.client is not None
244246
assert embedder.client.max_retries == 5
245247
assert embedder.client.timeout == 30.0
246248

247249
def test_warm_up_uses_timeout_and_max_retries_from_parameters(self):
248250
embedder = OpenAITextEmbedder(api_key=Secret.from_token("fake-api-key"), timeout=40.0, max_retries=1)
249251
embedder.warm_up()
252+
assert embedder.client is not None
250253
assert embedder.client.max_retries == 1
251254
assert embedder.client.timeout == 40.0
252255

@@ -255,6 +258,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
255258
monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
256259
embedder = OpenAITextEmbedder(api_key=Secret.from_token("fake-api-key"))
257260
embedder.warm_up()
261+
assert embedder.client is not None
258262
assert embedder.client.max_retries == 10
259263
assert embedder.client.timeout == 100.0
260264

@@ -266,6 +270,7 @@ def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
266270

267271
def test_sync_lifecycle(self, mock_openai_clients):
268272
sync_cls, _ = mock_openai_clients
273+
sync_client = sync_cls.return_value
269274
embedder = OpenAITextEmbedder()
270275
assert embedder.client is None
271276
assert embedder.async_client is None
@@ -275,19 +280,20 @@ def test_sync_lifecycle(self, mock_openai_clients):
275280
assert embedder.async_client is None
276281

277282
embedder.close()
278-
sync_cls.return_value.close.assert_called_once()
283+
sync_client.close.assert_called_once()
279284
assert embedder.client is None
280285

281286
async def test_async_lifecycle(self, mock_openai_clients):
282287
_, async_cls = mock_openai_clients
288+
async_client = async_cls.return_value
283289
embedder = OpenAITextEmbedder()
284290

285291
await embedder.warm_up_async()
286292
assert embedder.async_client is async_cls.return_value
287293
assert embedder.client is None
288294

289295
await embedder.close_async()
290-
async_cls.return_value.close.assert_awaited_once()
296+
async_client.close.assert_awaited_once()
291297
assert embedder.async_client is None
292298

293299
async def test_close_is_safe_without_warm_up(self, mock_openai_clients):

0 commit comments

Comments
 (0)