test: adopt MockChatGenerator in tests

test/components/agents/test_agent_hitl.py

@@ -8,9 +8,8 @@
 
 import pytest
 
-from haystack import component
 from haystack.components.agents import Agent
-from haystack.components.generators.chat import OpenAIChatGenerator
+from haystack.components.generators.chat import MockChatGenerator, OpenAIChatGenerator
 from haystack.dataclasses import ChatMessage, ToolCall
 from haystack.hooks.human_in_the_loop import (
     AlwaysAskPolicy,
@@ -21,7 +20,7 @@
     SimpleConsoleUI,
 )
 from haystack.hooks.human_in_the_loop.types import ConfirmationStrategy, ConfirmationUI
-from haystack.tools import Tool, Toolset, create_tool_from_function
+from haystack.tools import Tool, create_tool_from_function
 
 
 class MockUserInterface(ConfirmationUI):
@@ -234,13 +233,6 @@ async def test_run_async_blocking_confirmation_strategy_modify(self, tools):
         assert result["token_usage"]["total_tokens"] > 0
 
 
-@component
-class MockChatGenerator:
-    @component.output_types(replies=list[ChatMessage])
-    def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
-        return {"replies": [ChatMessage.from_assistant("Hello")]}
-
-
 def _producer() -> dict[str, str]:
     return {"value": "PRODUCED"}
 

test/components/agents/test_agent_hooks.py

@@ -2,30 +2,17 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Annotated, Any
+from typing import Annotated
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
-from haystack import component
 from haystack.components.agents import Agent
 from haystack.components.agents.state import State, replace_values
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.dataclasses import ChatMessage, ToolCall
 from haystack.hooks import hook
-from haystack.tools import Tool, Toolset, tool
-
-
-@component
-class MockChatGenerator:
-    @component.output_types(replies=list[ChatMessage])
-    def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
-        return {"replies": [ChatMessage.from_assistant("Hello")]}
-
-    @component.output_types(replies=list[ChatMessage])
-    async def run_async(
-        self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
-    ) -> dict[str, Any]:
-        return {"replies": [ChatMessage.from_assistant("Hello")]}
+from haystack.tools import tool
 
 
 @tool

test/components/extractors/image/test_llm_document_content_extractor.py

@@ -11,7 +11,7 @@
 from haystack import Document, Pipeline
 from haystack.components.converters.image.document_to_image import DocumentToImageContent
 from haystack.components.extractors.image import LLMDocumentContentExtractor
-from haystack.components.generators.chat import OpenAIChatGenerator
+from haystack.components.generators.chat import MockChatGenerator, OpenAIChatGenerator
 from haystack.components.writers import DocumentWriter
 from haystack.core.serialization import component_to_dict
 from haystack.dataclasses.chat_message import ChatMessage, ImageContent
@@ -191,12 +191,10 @@ def test_run_with_llm_success(self, mock_doc_to_image_run):
     @patch.object(DocumentToImageContent, "run")
     def test_run_plain_string_response_goes_to_content(self, mock_doc_to_image_run):
         """When LLM returns plain string (non-JSON), it is written to document content."""
-        mock_chat_generator = Mock(spec=OpenAIChatGenerator)
-        mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant(text="Plain text, not JSON")]}
         mock_doc_to_image_run.return_value = {
             "image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
         }
-        extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
+        extractor = LLMDocumentContentExtractor(chat_generator=MockChatGenerator("Plain text, not JSON"))
         docs = [Document(content="", meta={"file_path": "/path/to/image.pdf"})]
         result = extractor.run(documents=docs)
         assert len(result["documents"]) == 1
@@ -206,14 +204,10 @@ def test_run_plain_string_response_goes_to_content(self, mock_doc_to_image_run):
     @patch.object(DocumentToImageContent, "run")
     def test_run_valid_json_not_object_reports_error(self, mock_doc_to_image_run):
         """When LLM returns valid JSON that is not an object (e.g. array or primitive), report error."""
-        mock_chat_generator = Mock(spec=OpenAIChatGenerator)
-        mock_chat_generator.run.return_value = {
-            "replies": [ChatMessage.from_assistant(text='["array", "not", "object"]')]
-        }
         mock_doc_to_image_run.return_value = {
             "image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
         }
-        extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
+        extractor = LLMDocumentContentExtractor(chat_generator=MockChatGenerator('["array", "not", "object"]'))
         docs = [Document(content="", meta={"file_path": "/path/to/image.pdf"})]
         result = extractor.run(documents=docs)
         assert len(result["documents"]) == 0
@@ -224,16 +218,12 @@ def test_run_valid_json_not_object_reports_error(self, mock_doc_to_image_run):
     @patch.object(DocumentToImageContent, "run")
     def test_run_with_content_and_metadata_extraction(self, mock_doc_to_image_run):
         """When content mode gets JSON with document_content and other keys, other keys are merged into metadata."""
-        mock_chat_generator = Mock(spec=OpenAIChatGenerator)
-        mock_chat_generator.run.return_value = {
-            "replies": [
-                ChatMessage.from_assistant(text='{"document_content": "Main text", "title": "Doc Title", "page": "1"}')
-            ]
-        }
         mock_doc_to_image_run.return_value = {
             "image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
         }
-        extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
+        extractor = LLMDocumentContentExtractor(
+            chat_generator=MockChatGenerator('{"document_content": "Main text", "title": "Doc Title", "page": "1"}')
+        )
         docs = [Document(content="", meta={"file_path": "/path/to/image.pdf"})]
         result = extractor.run(documents=docs)
         assert len(result["documents"]) == 1
@@ -285,16 +275,14 @@ def test_run_with_llm_failure_raise_on_failure_true(self, mock_doc_to_image_run)
 
     @patch.object(DocumentToImageContent, "run")
     def test_run_removes_extraction_error_from_previous_runs(self, mock_doc_to_image_run):
-        mock_chat_generator = Mock(spec=OpenAIChatGenerator)
-        mock_chat_generator.run.return_value = {
-            "replies": [ChatMessage.from_assistant(text='{"document_content": "Successfully extracted content"}')]
-        }
         # Mock DocumentToImageContent to return valid image content
         mock_doc_to_image_run.return_value = {
             "image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
         }
 
-        extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
+        extractor = LLMDocumentContentExtractor(
+            chat_generator=MockChatGenerator('{"document_content": "Successfully extracted content"}')
+        )
 
         # Document with previous extraction error
         docs = [
@@ -356,16 +344,12 @@ def test_run_with_mixed_success_and_failure(self, mock_doc_to_image_run):
     @patch.object(DocumentToImageContent, "run")
     def test_run_json_multiple_keys_metadata_merged(self, mock_doc_to_image_run):
         """When LLM returns JSON with multiple keys and no document_content, all keys are merged into metadata."""
-        mock_chat_generator = Mock(spec=OpenAIChatGenerator)
-        mock_chat_generator.run.return_value = {
-            "replies": [
-                ChatMessage.from_assistant(text='{"title": "Sample Doc", "author": "Test", "document_type": "report"}')
-            ]
-        }
         mock_doc_to_image_run.return_value = {
             "image_contents": [ImageContent.from_file_path("./test/test_files/images/apple.jpg")]
         }
-        extractor = LLMDocumentContentExtractor(chat_generator=mock_chat_generator)
+        extractor = LLMDocumentContentExtractor(
+            chat_generator=MockChatGenerator('{"title": "Sample Doc", "author": "Test", "document_type": "report"}')
+        )
         original_content = "Original content"
         docs = [Document(content=original_content, meta={"file_path": "/path/to/image.pdf"})]
         result = extractor.run(documents=docs)

test/components/query/test_query_expander.py

@@ -8,6 +8,7 @@
 
 import pytest
 
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.components.generators.chat.openai import OpenAIChatGenerator
 from haystack.components.query.query_expander import DEFAULT_PROMPT_TEMPLATE, QueryExpander
 from haystack.dataclasses.chat_message import ChatMessage
@@ -104,12 +105,9 @@ def test_run_successful_expansion(self, mock_chat_generator):
         ]
         mock_chat_generator.run.assert_called_once()
 
-    def test_run_without_including_original(self, mock_chat_generator):
-        mock_chat_generator.run.return_value = {
-            "replies": [ChatMessage.from_assistant('{"queries": ["alt1", "alt2"]}')]
-        }
-
-        expander = QueryExpander(chat_generator=mock_chat_generator, include_original_query=False)
+    def test_run_without_including_original(self):
+        chat_generator = MockChatGenerator('{"queries": ["alt1", "alt2"]}')
+        expander = QueryExpander(chat_generator=chat_generator, include_original_query=False)
         result = expander.run("original")
 
         assert result["queries"] == ["alt1", "alt2"]
@@ -149,9 +147,8 @@ def test_run_generator_exception(self, mock_chat_generator):
         result = expander.run("test query")
         assert result["queries"] == ["test query"]
 
-    def test_run_invalid_json_response(self, mock_chat_generator):
-        mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("invalid json response")]}
-        expander = QueryExpander(chat_generator=mock_chat_generator)
+    def test_run_invalid_json_response(self):
+        expander = QueryExpander(chat_generator=MockChatGenerator("invalid json response"))
         result = expander.run("test query")
         assert result["queries"] == ["test query"]
 
@@ -196,20 +193,16 @@ def test_parse_expanded_queries_mixed_types(self, monkeypatch):
         queries = expander._parse_expanded_queries('{"queries": ["valid query", 123, "", "another valid"]}')
         assert queries == ["valid query", "another valid"]
 
-    def test_run_query_deduplication(self, mock_chat_generator):
-        mock_chat_generator.run.return_value = {
-            "replies": [ChatMessage.from_assistant('{"queries": ["original query", "alt1", "alt2"]}')]
-        }
-        expander = QueryExpander(chat_generator=mock_chat_generator, include_original_query=True)
+    def test_run_query_deduplication(self):
+        chat_generator = MockChatGenerator('{"queries": ["original query", "alt1", "alt2"]}')
+        expander = QueryExpander(chat_generator=chat_generator, include_original_query=True)
         result = expander.run("original query")
         assert result["queries"] == ["original query", "alt1", "alt2"]
         assert len(result["queries"]) == 3
 
-    def test_run_truncates_excess_queries(self, mock_chat_generator, caplog):
-        mock_chat_generator.run.return_value = {
-            "replies": [ChatMessage.from_assistant('{"queries": ["q1", "q2", "q3", "q4", "q5"]}')]
-        }
-        expander = QueryExpander(chat_generator=mock_chat_generator, n_expansions=3, include_original_query=False)
+    def test_run_truncates_excess_queries(self, caplog):
+        chat_generator = MockChatGenerator('{"queries": ["q1", "q2", "q3", "q4", "q5"]}')
+        expander = QueryExpander(chat_generator=chat_generator, n_expansions=3, include_original_query=False)
 
         with caplog.at_level(logging.WARNING):
             result = expander.run("test query")
@@ -244,14 +237,8 @@ def test_run_with_custom_template(self, mock_chat_generator):
         assert "Create 2 alternative search queries for: test query" in call_args
         assert "Return as JSON" in call_args
 
-    def test_component_output_types(self, mock_chat_generator, monkeypatch):
-        monkeypatch.setenv("OPENAI_API_KEY", "test-key-12345")
-        expander = QueryExpander()
-
-        mock_chat_generator.run.return_value = {
-            "replies": [ChatMessage.from_assistant('{"queries": ["test1", "test2"]}')]
-        }
-        expander.chat_generator = mock_chat_generator
+    def test_component_output_types(self):
+        expander = QueryExpander(chat_generator=MockChatGenerator('{"queries": ["test1", "test2"]}'))
 
         result = expander.run("test")
         assert "queries" in result