test: type-check test/components/embedders and add it to mypy config

pyproject.toml

@@ -149,7 +149,7 @@ all = 'pytest {args:test}'
 e2e = 'pytest {args:e2e}'
 
 # TODO We want to eventually type the whole test folder
-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/}"
+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/}"
 
 [project.urls]
 "CI: GitHub" = "https://github.com/deepset-ai/haystack/actions"

test/components/embedders/test_azure_document_embedder.py

@@ -216,6 +216,7 @@ def test_embed_batch_handles_exceptions_gracefully(self, caplog):
             embedding_separator=" | ",
         )
         embedder.warm_up()
+        assert embedder.client is not None
 
         fake_texts_to_embed = {"1": "text1", "2": "text2"}
 
@@ -237,6 +238,7 @@ def test_embed_batch_raises_exception_on_failure(self):
             raise_on_failure=True,
         )
         embedder.warm_up()
+        assert embedder.client is not None
         fake_texts_to_embed = {"1": "text1", "2": "text2"}
         with patch.object(
             embedder.client.embeddings,
@@ -303,6 +305,7 @@ def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
         monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
         embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 5
         assert embedder.client.timeout == 30.0
 
@@ -312,6 +315,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_parameters(self, monkeypatch)
             azure_endpoint="https://example-resource.azure.openai.com/", timeout=40.0, max_retries=1
         )
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 1
         assert embedder.client.timeout == 40.0
 
@@ -321,6 +325,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
         monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
         embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 10
         assert embedder.client.timeout == 100.0
 
@@ -334,6 +339,7 @@ def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
 
     def test_sync_lifecycle(self, mock_azure_clients):
         sync_cls, _ = mock_azure_clients
+        sync_client = sync_cls.return_value
         embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
         assert embedder.client is None
         assert embedder.async_client is None
@@ -343,19 +349,20 @@ def test_sync_lifecycle(self, mock_azure_clients):
         assert embedder.async_client is None
 
         embedder.close()
-        sync_cls.return_value.close.assert_called_once()
+        sync_client.close.assert_called_once()
         assert embedder.client is None
 
     async def test_async_lifecycle(self, mock_azure_clients):
         _, async_cls = mock_azure_clients
+        async_client = async_cls.return_value
         embedder = AzureOpenAIDocumentEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
 
         await embedder.warm_up_async()
         assert embedder.async_client is async_cls.return_value
         assert embedder.client is None
 
         await embedder.close_async()
-        async_cls.return_value.close.assert_awaited_once()
+        async_client.close.assert_awaited_once()
         assert embedder.async_client is None
 
     async def test_close_is_safe_without_warm_up(self, mock_azure_clients):

test/components/embedders/test_azure_text_embedder.py

@@ -216,6 +216,7 @@ def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
         monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake-api-key")
         embedder = AzureOpenAITextEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 5
         assert embedder.client.timeout == 30.0
 
@@ -225,6 +226,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_parameters(self, monkeypatch)
             azure_endpoint="https://example-resource.azure.openai.com/", timeout=40.0, max_retries=1
         )
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 1
         assert embedder.client.timeout == 40.0
 
@@ -234,6 +236,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
         monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
         embedder = AzureOpenAITextEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 10
         assert embedder.client.timeout == 100.0
 
@@ -247,6 +250,7 @@ def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
 
     def test_sync_lifecycle(self, mock_azure_clients):
         sync_cls, _ = mock_azure_clients
+        sync_client = sync_cls.return_value
         embedder = AzureOpenAITextEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
         assert embedder.client is None
         assert embedder.async_client is None
@@ -256,19 +260,20 @@ def test_sync_lifecycle(self, mock_azure_clients):
         assert embedder.async_client is None
 
         embedder.close()
-        sync_cls.return_value.close.assert_called_once()
+        sync_client.close.assert_called_once()
         assert embedder.client is None
 
     async def test_async_lifecycle(self, mock_azure_clients):
         _, async_cls = mock_azure_clients
+        async_client = async_cls.return_value
         embedder = AzureOpenAITextEmbedder(azure_endpoint="https://example-resource.azure.openai.com/")
 
         await embedder.warm_up_async()
         assert embedder.async_client is async_cls.return_value
         assert embedder.client is None
 
         await embedder.close_async()
-        async_cls.return_value.close.assert_awaited_once()
+        async_client.close.assert_awaited_once()
         assert embedder.async_client is None
 
     async def test_close_is_safe_without_warm_up(self, mock_azure_clients):

test/components/embedders/test_mock_text_embedder.py

@@ -60,7 +60,8 @@ def test_embedding_fn(self):
         assert MockTextEmbedder(embedding_fn=_ones).run("hello")["embedding"] == [1.0, 1.0, 1.0]
 
     def test_embedding_fn_invalid_return_raises(self):
-        embedder = MockTextEmbedder(embedding_fn=lambda text: "not a vector")
+        # embedding_fn deliberately returns a non-vector to exercise the runtime type check
+        embedder = MockTextEmbedder(embedding_fn=lambda text: "not a vector")  # type: ignore[arg-type, return-value]
         with pytest.raises(TypeError, match="must be a sequence of numbers"):
             embedder.run("hello")
 
@@ -79,8 +80,9 @@ def test_meta(self):
         assert custom["extra"] == "value"
 
     def test_run_rejects_non_string(self):
+        # a non-string input is passed on purpose to exercise the runtime type check
         with pytest.raises(TypeError, match="expects a string"):
-            MockTextEmbedder().run(["not", "a", "string"])
+            MockTextEmbedder().run(["not", "a", "string"])  # type: ignore[arg-type]
 
     async def test_run_async(self):
         embedder = MockTextEmbedder(dimension=8)

test/components/embedders/test_openai_document_embedder.py

@@ -200,15 +200,15 @@ def test_run_wrong_input_format(self):
         list_integers_input = [1, 2, 3]
 
         with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"):
-            embedder.run(documents=string_input)
+            embedder.run(documents=string_input)  # type: ignore[arg-type]
 
         with pytest.raises(TypeError, match="OpenAIDocumentEmbedder expects a list of Documents as input"):
-            embedder.run(documents=list_integers_input)
+            embedder.run(documents=list_integers_input)  # type: ignore[arg-type]
 
     def test_run_on_empty_list(self):
         embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake-api-key"))
 
-        empty_list_input = []
+        empty_list_input: list[Document] = []
         result = embedder.run(documents=empty_list_input)
 
         assert result["documents"] is not None
@@ -217,6 +217,7 @@ def test_run_on_empty_list(self):
     def test_embed_batch_handles_exceptions_gracefully(self, caplog):
         embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake_api_key"))
         embedder.warm_up()
+        assert embedder.client is not None
         fake_texts_to_embed = {"1": "text1", "2": "text2"}
         with patch.object(
             embedder.client.embeddings,
@@ -231,6 +232,7 @@ def test_embed_batch_handles_exceptions_gracefully(self, caplog):
     def test_run_handles_exceptions_gracefully(self, caplog):
         embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake_api_key"), batch_size=1)
         embedder.warm_up()
+        assert embedder.client is not None
         docs = [
             Document(content="I love cheese", meta={"topic": "Cuisine"}),
             Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}),
@@ -260,6 +262,7 @@ def test_run_handles_exceptions_gracefully(self, caplog):
     def test_embed_batch_raises_exception_on_failure(self):
         embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake_api_key"), raise_on_failure=True)
         embedder.warm_up()
+        assert embedder.client is not None
         fake_texts_to_embed = {"1": "text1", "2": "text2"}
         with patch.object(
             embedder.client.embeddings,
@@ -354,12 +357,14 @@ def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
         monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
         embedder = OpenAIDocumentEmbedder()
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 5
         assert embedder.client.timeout == 30.0
 
     def test_warm_up_uses_timeout_and_max_retries_from_parameters(self):
         embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake-api-key"), timeout=40.0, max_retries=1)
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 1
         assert embedder.client.timeout == 40.0
 
@@ -368,6 +373,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
         monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
         embedder = OpenAIDocumentEmbedder(api_key=Secret.from_token("fake-api-key"))
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 10
         assert embedder.client.timeout == 100.0
 
@@ -379,6 +385,7 @@ def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
 
     def test_sync_lifecycle(self, mock_openai_clients):
         sync_cls, _ = mock_openai_clients
+        sync_client = sync_cls.return_value
         embedder = OpenAIDocumentEmbedder()
         assert embedder.client is None
         assert embedder.async_client is None
@@ -388,19 +395,20 @@ def test_sync_lifecycle(self, mock_openai_clients):
         assert embedder.async_client is None
 
         embedder.close()
-        sync_cls.return_value.close.assert_called_once()
+        sync_client.close.assert_called_once()
         assert embedder.client is None
 
     async def test_async_lifecycle(self, mock_openai_clients):
         _, async_cls = mock_openai_clients
+        async_client = async_cls.return_value
         embedder = OpenAIDocumentEmbedder()
 
         await embedder.warm_up_async()
         assert embedder.async_client is async_cls.return_value
         assert embedder.client is None
 
         await embedder.close_async()
-        async_cls.return_value.close.assert_awaited_once()
+        async_client.close.assert_awaited_once()
         assert embedder.async_client is None
 
     async def test_close_is_safe_without_warm_up(self, mock_openai_clients):

test/components/embedders/test_openai_text_embedder.py

@@ -8,6 +8,7 @@
 
 import pytest
 from openai.types import CreateEmbeddingResponse, Embedding
+from openai.types.create_embedding_response import Usage
 
 import haystack.components.embedders.openai_text_embedder as openai_text_embedder_module
 from haystack.components.embedders.openai_text_embedder import OpenAITextEmbedder
@@ -169,7 +170,7 @@ def test_prepare_output(self, monkeypatch):
             data=[Embedding(embedding=[0.1, 0.2, 0.3], index=0, object="embedding")],
             model="text-embedding-ada-002",
             object="list",
-            usage={"prompt_tokens": 6, "total_tokens": 6},
+            usage=Usage(prompt_tokens=6, total_tokens=6),
         )
 
         embedder = OpenAITextEmbedder()
@@ -185,7 +186,7 @@ def test_run_wrong_input_format(self):
         list_integers_input = [1, 2, 3]
 
         with pytest.raises(TypeError, match="OpenAITextEmbedder expects a string as an input"):
-            embedder.run(text=list_integers_input)
+            embedder.run(text=list_integers_input)  # type: ignore[arg-type]
 
     @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")
     @pytest.mark.integration
@@ -241,12 +242,14 @@ def test_warm_up_uses_default_timeout_and_max_retries(self, monkeypatch):
         monkeypatch.setenv("OPENAI_API_KEY", "fake-api-key")
         embedder = OpenAITextEmbedder()
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 5
         assert embedder.client.timeout == 30.0
 
     def test_warm_up_uses_timeout_and_max_retries_from_parameters(self):
         embedder = OpenAITextEmbedder(api_key=Secret.from_token("fake-api-key"), timeout=40.0, max_retries=1)
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 1
         assert embedder.client.timeout == 40.0
 
@@ -255,6 +258,7 @@ def test_warm_up_uses_timeout_and_max_retries_from_env_vars(self, monkeypatch):
         monkeypatch.setenv("OPENAI_MAX_RETRIES", "10")
         embedder = OpenAITextEmbedder(api_key=Secret.from_token("fake-api-key"))
         embedder.warm_up()
+        assert embedder.client is not None
         assert embedder.client.max_retries == 10
         assert embedder.client.timeout == 100.0
 
@@ -266,6 +270,7 @@ def test_key_resolved_at_warm_up_not_init(self, monkeypatch):
 
     def test_sync_lifecycle(self, mock_openai_clients):
         sync_cls, _ = mock_openai_clients
+        sync_client = sync_cls.return_value
         embedder = OpenAITextEmbedder()
         assert embedder.client is None
         assert embedder.async_client is None
@@ -275,19 +280,20 @@ def test_sync_lifecycle(self, mock_openai_clients):
         assert embedder.async_client is None
 
         embedder.close()
-        sync_cls.return_value.close.assert_called_once()
+        sync_client.close.assert_called_once()
         assert embedder.client is None
 
     async def test_async_lifecycle(self, mock_openai_clients):
         _, async_cls = mock_openai_clients
+        async_client = async_cls.return_value
         embedder = OpenAITextEmbedder()
 
         await embedder.warm_up_async()
         assert embedder.async_client is async_cls.return_value
         assert embedder.client is None
 
         await embedder.close_async()
-        async_cls.return_value.close.assert_awaited_once()
+        async_client.close.assert_awaited_once()
         assert embedder.async_client is None
 
     async def test_close_is_safe_without_warm_up(self, mock_openai_clients):