diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py
index e0517620b64..35df8ccc243 100644
--- a/test/components/agents/test_agent.py
+++ b/test/components/agents/test_agent.py
@@ -21,6 +21,7 @@
 from haystack.components.agents.tool_calling import _run_tool
 from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
 from haystack.components.builders.prompt_builder import PromptBuilder
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.components.generators.chat.openai import OpenAIChatGenerator
 from haystack.components.joiners.branch import BranchJoiner
 from haystack.components.joiners.list_joiner import ListJoiner
@@ -101,7 +102,7 @@ def component_tool():
 @pytest.fixture
 def make_agent(weather_tool):
     def _factory(**kwargs):
-        return Agent(chat_generator=MockChatGenerator(), tools=[weather_tool], **kwargs)
+        return Agent(chat_generator=MockChatGenerator("Hello"), tools=[weather_tool], **kwargs)
 
     return _factory
 
@@ -177,46 +178,37 @@ def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None =
 
 
 @component
-class MockChatGenerator:
-    def to_dict(self) -> dict[str, Any]:
-        return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}}
+class ToolAssertingChatGenerator:
+    """Asserts the Agent forwards the expected tools, then drives one tool call before a plain reply."""
 
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "MockChatGenerator":
-        return cls()
+    def __init__(self, expected_tools):
+        self.expected_tools = expected_tools
+        self.tool_invoked = False
 
     @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 run_async")]}
+        assert tools == self.expected_tools
+        tool_message = ChatMessage.from_assistant(
+            tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
+        )
+        message = tool_message if not self.tool_invoked else ChatMessage.from_assistant("Hello")
+        self.tool_invoked = True
+        return {"replies": [message]}
 
 
-@component
-class ParallelToolCallingChatGenerator:
+def _parallel_tool_calling_generator() -> MockChatGenerator:
     """Requests two `weather_tool` calls on the first turn, then returns a plain reply so the agent loop exits."""
-
-    tool_invoked = False
-
-    @component.output_types(replies=list[ChatMessage])
-    def run(self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs) -> dict[str, Any]:
-        if self.tool_invoked:
-            return {"replies": [ChatMessage.from_assistant("done")]}
-        self.tool_invoked = True
-        return {
-            "replies": [
-                ChatMessage.from_assistant(
-                    tool_calls=[
-                        ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"}),
-                        ToolCall(tool_name="weather_tool", arguments={"location": "Paris"}),
-                    ]
-                )
-            ]
-        }
+    return MockChatGenerator(
+        [
+            ChatMessage.from_assistant(
+                tool_calls=[
+                    ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"}),
+                    ToolCall(tool_name="weather_tool", arguments={"location": "Paris"}),
+                ]
+            ),
+            "done",
+        ]
+    )
 
 
 class TestAgent:
@@ -791,21 +783,13 @@ def test_exceed_max_steps(self, monkeypatch, weather_tool, caplog):
             assert "Agent reached maximum agent steps" in caplog.text
             assert result["step_count"] == 0
 
-    def test_exit_condition_exits(self, monkeypatch, weather_tool):
-        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
-        generator = OpenAIChatGenerator()
-
-        # Mock messages where the exit condition appears in the second message
-        mock_messages = [
-            ChatMessage.from_assistant(
-                tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
-            )
-        ]
-
-        agent = Agent(chat_generator=generator, tools=[weather_tool], exit_conditions=["weather_tool"])
-
-        # Patch agent.chat_generator.run to return mock_messages
-        agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
+    def test_exit_condition_exits(self, weather_tool):
+        tool_call_message = ChatMessage.from_assistant(
+            tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
+        )
+        agent = Agent(
+            chat_generator=MockChatGenerator(tool_call_message), tools=[weather_tool], exit_conditions=["weather_tool"]
+        )
 
         result = agent.run([ChatMessage.from_user("Hello")])
 
@@ -820,23 +804,20 @@ def test_exit_condition_exits(self, monkeypatch, weather_tool):
         assert isinstance(result["last_message"], ChatMessage)
         assert result["messages"][-1] == result["last_message"]
 
-    def test_exit_condition_on_tool_provided_at_runtime(self, monkeypatch, weather_tool):
+    def test_exit_condition_on_tool_provided_at_runtime(self, weather_tool):
         """An exit condition naming a tool absent at init still triggers once that tool is provided at runtime."""
-        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
-
-        # weather_tool is NOT among the init tools, but it is named as an exit condition.
+        # weather_tool is NOT among the init tools, but it is named as an exit condition. The model calls
+        # weather_tool, which is supplied only at runtime.
+        tool_call_message = ChatMessage.from_assistant(
+            tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
+        )
         agent = Agent(
-            chat_generator=OpenAIChatGenerator(), tools=[], exit_conditions=["weather_tool"], max_agent_steps=5
+            chat_generator=MockChatGenerator(tool_call_message),
+            tools=[],
+            exit_conditions=["weather_tool"],
+            max_agent_steps=5,
         )
 
-        # The model calls weather_tool, which is supplied only at runtime.
-        mock_messages = [
-            ChatMessage.from_assistant(
-                tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
-            )
-        ]
-        agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
-
         result = agent.run([ChatMessage.from_user("What's the weather in Berlin?")], tools=[weather_tool])
 
         # The agent exits right after the exit-condition tool runs (single step), not at max_agent_steps.
@@ -911,17 +892,8 @@ def test_check_exit_conditions_parallel_calls_error_only_on_non_exit_tool(
 
         assert agent._check_exit_conditions(llm_messages, tool_messages) is True
 
-    def test_agent_with_no_tools(self, monkeypatch):
-        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
-        generator = OpenAIChatGenerator()
-
-        # Mock messages where the exit condition appears in the second message
-        mock_messages = [ChatMessage.from_assistant("Berlin")]
-
-        agent = Agent(chat_generator=generator, tools=[], max_agent_steps=3)
-
-        # Patch agent.chat_generator.run to return mock_messages
-        agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
+    def test_agent_with_no_tools(self):
+        agent = Agent(chat_generator=MockChatGenerator("Berlin"), tools=[], max_agent_steps=3)
 
         response = agent.run([ChatMessage.from_user("What is the capital of Germany?")])
 
@@ -943,23 +915,7 @@ def test_run_with_system_prompt(self, weather_tool):
         assert response["messages"][0].text == "This is a system prompt."
 
     def test_run_with_tools_run_param(self, weather_tool: Tool, component_tool: Tool, monkeypatch):
-        @component
-        class MockChatGenerator:
-            tool_invoked = False
-
-            @component.output_types(replies=list[ChatMessage])
-            def run(
-                self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
-            ) -> dict[str, Any]:
-                assert tools == [weather_tool]
-                tool_message = ChatMessage.from_assistant(
-                    tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
-                )
-                message = tool_message if not self.tool_invoked else ChatMessage.from_assistant("Hello")
-                self.tool_invoked = True
-                return {"replies": [message]}
-
-        chat_generator = MockChatGenerator()
+        chat_generator = ToolAssertingChatGenerator(expected_tools=[weather_tool])
         agent = Agent(
             chat_generator=chat_generator,
             tools=[component_tool],
@@ -975,23 +931,7 @@ def run(
         assert run_tool_mock.call_args.kwargs["enable_streaming_callback_passthrough"] is True
 
     def test_run_with_tools_run_param_for_tool_selection(self, weather_tool: Tool, component_tool: Tool, monkeypatch):
-        @component
-        class MockChatGenerator:
-            tool_invoked = False
-
-            @component.output_types(replies=list[ChatMessage])
-            def run(
-                self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
-            ) -> dict[str, Any]:
-                assert tools == [weather_tool]
-                tool_message = ChatMessage.from_assistant(
-                    tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
-                )
-                message = tool_message if not self.tool_invoked else ChatMessage.from_assistant("Hello")
-                self.tool_invoked = True
-                return {"replies": [message]}
-
-        chat_generator = MockChatGenerator()
+        chat_generator = ToolAssertingChatGenerator(expected_tools=[weather_tool])
         agent = Agent(
             chat_generator=chat_generator,
             tools=[weather_tool, component_tool],
@@ -1055,7 +995,7 @@ def test_run(self, weather_tool):
 
     @pytest.mark.asyncio
     async def test_generation_kwargs(self):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
 
         agent = Agent(chat_generator=chat_generator)
 
@@ -1073,7 +1013,7 @@ async def test_generation_kwargs(self):
 
     @pytest.mark.asyncio
     async def test_run_async_uses_chat_generator_run_async_when_available(self, weather_tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
 
         chat_generator.run_async = AsyncMock(
@@ -1163,7 +1103,7 @@ async def test_does_not_exit_on_empty_assistant_message_async(self, monkeypatch,
 
     @pytest.mark.asyncio
     async def test_run_async_with_async_streaming_callback(self, weather_tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=async_streaming_callback)
 
         # This should not raise any exception
@@ -1171,10 +1111,10 @@ async def test_run_async_with_async_streaming_callback(self, weather_tool):
 
         assert "messages" in result
         assert len(result["messages"]) == 2
-        assert result["messages"][1].text == "Hello from run_async"
+        assert result["messages"][1].text == "Hello"
 
     def test_run_with_async_streaming_callback_fails(self, weather_tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=async_streaming_callback)
 
         with pytest.raises(ValueError, match="The init callback cannot be a coroutine"):
@@ -1182,7 +1122,7 @@ def test_run_with_async_streaming_callback_fails(self, weather_tool):
 
     @pytest.mark.asyncio
     async def test_run_async_with_sync_streaming_callback_warns(self, weather_tool, caplog):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=sync_streaming_callback)
 
         with caplog.at_level(logging.WARNING):
@@ -1200,10 +1140,8 @@ def test_reserved_state_schema_keys_raise(self, monkeypatch, weather_tool):
                     chat_generator=OpenAIChatGenerator(), tools=[weather_tool], state_schema={reserved: {"type": int}}
                 )
 
-    def test_run_populates_token_usage_and_tool_call_counts(self, monkeypatch, weather_tool, component_tool):
+    def test_run_populates_token_usage_and_tool_call_counts(self, weather_tool, component_tool):
         """A multi-step run aggregates step_count, token_usage (incl. nested details), and tool_call_counts."""
-        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
-        agent = Agent(chat_generator=OpenAIChatGenerator(), tools=[weather_tool, component_tool])
         # Step 1: two parallel tool calls + usage with nested detail dicts.
         # Step 2: one more weather_tool call + flat usage.
         # Step 3: final text answer + usage.
@@ -1238,8 +1176,9 @@ def test_run_populates_token_usage_and_tool_call_counts(self, monkeypatch, weath
                 },
             )
         ]
-        agent.chat_generator.run = MagicMock(
-            side_effect=[{"replies": first_step}, {"replies": second_step}, {"replies": third_step}]
+        agent = Agent(
+            chat_generator=MockChatGenerator(first_step + second_step + third_step),
+            tools=[weather_tool, component_tool],
         )
 
         result = agent.run([ChatMessage.from_user("Hi")])
@@ -1253,9 +1192,7 @@ def test_run_populates_token_usage_and_tool_call_counts(self, monkeypatch, weath
         }
 
     @pytest.mark.asyncio
-    async def test_run_async_populates_token_usage_and_tool_call_counts(self, monkeypatch, weather_tool):
-        monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
-        agent = Agent(chat_generator=OpenAIChatGenerator(), tools=[weather_tool])
+    async def test_run_async_populates_token_usage_and_tool_call_counts(self, weather_tool):
         first_step = [
             _assistant_with_usage(
                 tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})],
@@ -1265,7 +1202,7 @@ async def test_run_async_populates_token_usage_and_tool_call_counts(self, monkey
         second_step = [
             _assistant_with_usage("Done.", usage={"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4})
         ]
-        agent.chat_generator.run_async = AsyncMock(side_effect=[{"replies": first_step}, {"replies": second_step}])
+        agent = Agent(chat_generator=MockChatGenerator(first_step + second_step), tools=[weather_tool])
 
         result = await agent.run_async([ChatMessage.from_user("Hi")])
         assert result["step_count"] == 2
@@ -1274,9 +1211,10 @@ async def test_run_async_populates_token_usage_and_tool_call_counts(self, monkey
 
     def test_metadata_outputs_show_defaults_when_no_data(self, weather_tool):
         """`token_usage` stays empty and `tool_call_counts` reports zero for every tool when nothing happens."""
-        agent = Agent(chat_generator=MockChatGenerator(), tools=[weather_tool])
+        # A text-only reply whose `usage` meta is empty leaves `token_usage` empty after aggregation.
+        chat_generator = MockChatGenerator(ChatMessage.from_assistant("Hello", meta={"usage": {}}))
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
         result = agent.run([ChatMessage.from_user("Hi")])
-        # MockChatGenerator returns a text-only reply with no `usage` meta and no tool calls.
         assert result["step_count"] == 1
         assert result["token_usage"] == {}
         assert result["tool_call_counts"] == {"weather_tool": 0}
@@ -1348,7 +1286,7 @@ def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool):
             elif hasattr(record, "tag_name") and spans:
                 spans[-1][1][record.tag_name] = record.tag_value
 
-        # Keep only the agent's own spans. With the MockChatGenerator returning no tool calls, the inner
+        # Keep only the agent's own spans. With the chat generator returning no tool calls, the inner
         # `haystack.agent.step.tool` span never fires - the loop exits after the LLM call.
         agent_spans = [(op, tags) for op, tags in spans if op.startswith("haystack.agent")]
 
@@ -1415,26 +1353,15 @@ def test_agent_tracing_span_run_reflects_runtime_tools(self, caplog, monkeypatch
         tracing.disable_tracing()
 
     def test_agent_tracing_span_run_with_tool_call(self, caplog, monkeypatch, weather_tool):
-        @component
-        class ToolCallingChatGenerator:
-            tool_invoked = False
-
-            @component.output_types(replies=list[ChatMessage])
-            def run(
-                self, messages: list[ChatMessage], tools: list[Tool] | Toolset | None = None, **kwargs
-            ) -> dict[str, Any]:
-                if self.tool_invoked:
-                    return {"replies": [ChatMessage.from_assistant("done")]}
-                self.tool_invoked = True
-                return {
-                    "replies": [
-                        ChatMessage.from_assistant(
-                            tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
-                        )
-                    ]
-                }
-
-        agent = Agent(chat_generator=ToolCallingChatGenerator(), tools=[weather_tool])
+        chat_generator = MockChatGenerator(
+            [
+                ChatMessage.from_assistant(
+                    tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]
+                ),
+                "done",
+            ]
+        )
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
 
         tracing.tracer.is_content_tracing_enabled = True
         tracing.enable_tracing(LoggingTracer())
@@ -1478,7 +1405,7 @@ def run(
 
     def test_agent_tracing_span_run_with_parallel_tool_calls(self, caplog, monkeypatch, weather_tool):
         """Each tool call in a step gets its own `haystack.agent.step.tool` span instead of one grouped span."""
-        agent = Agent(chat_generator=ParallelToolCallingChatGenerator(), tools=[weather_tool])
+        agent = Agent(chat_generator=_parallel_tool_calling_generator(), tools=[weather_tool])
 
         tracing.tracer.is_content_tracing_enabled = True
         tracing.enable_tracing(LoggingTracer())
@@ -1515,7 +1442,7 @@ def test_agent_tracing_span_run_with_parallel_tool_calls(self, caplog, monkeypat
 
     @pytest.mark.asyncio
     async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
 
         tracing.tracer.is_content_tracing_enabled = True
@@ -1544,7 +1471,7 @@ async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_t
         )
         assert (
             llm_tags["haystack.agent.step.llm.output"]
-            == '{"replies": [{"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}]}'  # noqa: E501
+            == '{"replies": [{"role": "assistant", "meta": {"model": "mock-model", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, "name": null, "content": [{"text": "Hello"}]}]}'  # noqa: E501
         )
 
         _, step_tags = agent_spans[1]
@@ -1557,7 +1484,7 @@ async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_t
             "haystack.agent.exit_conditions": '["text"]',
             "haystack.agent.state_schema": '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}, "step_count": {"type": "int", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "token_usage": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tool_call_counts": {"type": "dict[str, int]", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "continue_run": {"type": "bool", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "tools": {"type": "list", "handler": "haystack.components.agents.state.state_utils.replace_values"}, "hook_context": {"type": "dict[str, typing.Any]", "handler": "haystack.components.agents.state.state_utils.replace_values"}}',  # noqa: E501
             "haystack.agent.input": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}',  # noqa: E501
-            "haystack.agent.output": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}], "step_count": 1, "token_usage": {}, "tool_call_counts": {"weather_tool": 0}, "last_message": {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}}',  # noqa: E501
+            "haystack.agent.output": '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {"model": "mock-model", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, "name": null, "content": [{"text": "Hello"}]}], "step_count": 1, "token_usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, "tool_call_counts": {"weather_tool": 0}, "last_message": {"role": "assistant", "meta": {"model": "mock-model", "index": 0, "finish_reason": "stop", "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, "name": null, "content": [{"text": "Hello"}]}}',  # noqa: E501
             "haystack.agent.steps_taken": 1,
         }
 
@@ -1568,7 +1495,7 @@ async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_t
     @pytest.mark.asyncio
     async def test_agent_tracing_span_async_run_with_parallel_tool_calls(self, caplog, monkeypatch, weather_tool):
         """The async path also emits one `haystack.agent.step.tool` span per tool call."""
-        agent = Agent(chat_generator=ParallelToolCallingChatGenerator(), tools=[weather_tool])
+        agent = Agent(chat_generator=_parallel_tool_calling_generator(), tools=[weather_tool])
 
         tracing.tracer.is_content_tracing_enabled = True
         tracing.enable_tracing(LoggingTracer())
@@ -1756,13 +1683,13 @@ def test_mixed_standalone_tools_and_toolsets(self, weather_tool: Tool, component
 
 class TestAgentToolSelection:
     def test_tool_selection_new_tool(self, weather_tool: Tool, component_tool: Tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.")
         result = agent._select_tools([component_tool])
         assert result == [component_tool]
 
     def test_tool_selection_existing_tools(self, weather_tool: Tool, component_tool: Tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(
             chat_generator=chat_generator,
             tools=[weather_tool, component_tool],
@@ -1772,7 +1699,7 @@ def test_tool_selection_existing_tools(self, weather_tool: Tool, component_tool:
         assert result == [weather_tool, component_tool]
 
     def test_tool_selection_invalid_type(self, weather_tool: Tool, component_tool: Tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(
             chat_generator=chat_generator,
             tools=[weather_tool, component_tool],
@@ -1790,7 +1717,7 @@ def test_tool_selection_invalid_type(self, weather_tool: Tool, component_tool: T
 
     def test_tool_selection_with_list_of_toolsets(self, weather_tool: Tool, component_tool: Tool):
         """Test that list of Toolsets and Tools can be passed to agent."""
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         toolset1 = Toolset([weather_tool])
         standalone_tool = Tool(
             name="standalone",
@@ -2054,7 +1981,7 @@ def test_rag_pipeline_user_prompt_init_only(self, make_rag_pipeline):
         assert "Documents:" in rendered
 
     def test_rag_pipeline_messages_plus_user_prompt(self, document_store_with_docs, weather_tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
 
         agent = Agent(
             chat_generator=chat_generator,
@@ -2142,7 +2069,7 @@ class AttachmentsBuilder:
             def run(self, processed_files: list[str]) -> dict:
                 return {"prompt": [ChatMessage.from_user(f"Files: {processed_files}")]}
 
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(
             chat_generator=chat_generator,
             tools=[weather_tool],
@@ -2209,7 +2136,7 @@ def tracking_warm_up():
 
     def test_warm_up_single_tool(self):
         tool = self._make_tracking_tool()
-        agent = Agent(chat_generator=MockChatGenerator(), tools=[tool])
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool])
 
         assert not tool.was_warmed_up
         agent.warm_up()
@@ -2218,7 +2145,7 @@ def test_warm_up_single_tool(self):
     def test_warm_up_multiple_tools(self):
         tool1 = self._make_tracking_tool("tool1")
         tool2 = self._make_tracking_tool("tool2")
-        agent = Agent(chat_generator=MockChatGenerator(), tools=[tool1, tool2])
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool1, tool2])
 
         assert not tool1.was_warmed_up
         assert not tool2.was_warmed_up
@@ -2229,7 +2156,7 @@ def test_warm_up_multiple_tools(self):
     def test_warm_up_toolset(self):
         inner_tool = self._make_tracking_tool()
         toolset = self._make_tracking_toolset([inner_tool])
-        agent = Agent(chat_generator=MockChatGenerator(), tools=toolset)
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=toolset)
 
         assert not toolset.was_warmed_up
         agent.warm_up()
@@ -2241,7 +2168,7 @@ def test_warm_up_mixed_toolsets(self):
         tool2 = self._make_tracking_tool("tool2")
         toolset2 = self._make_tracking_toolset([tool2])
 
-        agent = Agent(chat_generator=MockChatGenerator(), tools=toolset1 + toolset2)
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=toolset1 + toolset2)
 
         assert not toolset1.was_warmed_up
         assert not toolset2.was_warmed_up
@@ -2257,7 +2184,7 @@ def test_warm_up_mixed_list_of_tools_and_toolsets(self):
         tool4 = self._make_tracking_tool("toolset_tool2")
         toolset2 = self._make_tracking_toolset([tool4])
 
-        agent = Agent(chat_generator=MockChatGenerator(), tools=[tool1, toolset1, tool2, toolset2])
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool1, toolset1, tool2, toolset2])
 
         assert not tool1.was_warmed_up
         assert not tool2.was_warmed_up
@@ -2285,7 +2212,7 @@ def counting_warm_up():
 
         tool.warm_up = counting_warm_up
 
-        agent = Agent(chat_generator=MockChatGenerator(), tools=[tool])
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[tool])
         agent.warm_up()
         agent.warm_up()
         agent.warm_up()
@@ -2318,7 +2245,7 @@ def warm_up(self):
                     self._connected = True
 
         mcp_toolset = MockMCPToolset()
-        agent = Agent(chat_generator=MockChatGenerator(), tools=mcp_toolset)
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=mcp_toolset)
         assert mcp_toolset.tools == [placeholder_tool]
         agent.warm_up()
         assert mcp_toolset.tools == [actual_tool]
@@ -2377,7 +2304,7 @@ def run(self, messages: list[ChatMessage], tools: Toolset | None = None, **kwarg
     def test_run_warms_up_per_run_toolset(self):
         """Per-run tools passed to run() are not covered by Agent.warm_up() and must be warmed up at run time."""
         init_tool = self._make_tracking_tool("init_tool")
-        agent = Agent(chat_generator=MockChatGenerator(), tools=Toolset([init_tool]))
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=Toolset([init_tool]))
         agent.warm_up()
 
         per_run_tool = self._make_tracking_tool("per_run_tool")
@@ -2393,7 +2320,7 @@ def test_run_warms_up_per_run_toolset(self):
     def test_run_warms_up_per_run_list_of_tools_and_toolsets(self):
         """A per-run list of Tools and Toolsets must be warmed up at run time."""
         init_tool = self._make_tracking_tool("init_tool")
-        agent = Agent(chat_generator=MockChatGenerator(), tools=[init_tool])
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=[init_tool])
         agent.warm_up()
 
         per_run_tool = self._make_tracking_tool("per_run_tool")
@@ -2410,7 +2337,7 @@ def test_run_warms_up_per_run_list_of_tools_and_toolsets(self):
     async def test_run_async_warms_up_per_run_toolset(self):
         """The async run path must also warm up per-run tools."""
         init_tool = self._make_tracking_tool("init_tool")
-        agent = Agent(chat_generator=MockChatGenerator(), tools=Toolset([init_tool]))
+        agent = Agent(chat_generator=MockChatGenerator("Hello"), tools=Toolset([init_tool]))
         agent.warm_up()
 
         per_run_tool = self._make_tracking_tool("per_run_tool")
@@ -2424,7 +2351,7 @@ async def test_run_async_warms_up_per_run_toolset(self):
 
 class TestComponentLifecycle:
     def test_warm_up_delegates_to_chat_generator(self, weather_tool):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         chat_generator.warm_up = MagicMock()
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool], system_prompt="This is a system prompt.")
 
@@ -2434,11 +2361,12 @@ def test_warm_up_delegates_to_chat_generator(self, weather_tool):
         chat_generator.warm_up.reset_mock()
         agent.run([ChatMessage.from_user("What is the weather in Berlin?")])
         assert agent._tools_warmed_up is True
-        chat_generator.warm_up.assert_called_once()
+        # warm_up runs twice here: the Agent delegates to the generator, and the generator's own run() self-warms
+        assert chat_generator.warm_up.call_count == 2
 
     @pytest.mark.asyncio
     async def test_warm_up_async_delegates_to_chat_generator(self):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         chat_generator.warm_up_async = AsyncMock()
         chat_generator.warm_up = MagicMock()
         agent = Agent(chat_generator=chat_generator, tools=[])
@@ -2455,7 +2383,7 @@ async def test_warm_up_async_falls_back_to_sync_warm_up(self):
         chat_generator.warm_up.assert_called_once()
 
     def test_close_delegates_to_chat_generator(self):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         chat_generator.close = MagicMock()
         agent = Agent(chat_generator=chat_generator, tools=[])
         agent.close()
@@ -2463,7 +2391,7 @@ def test_close_delegates_to_chat_generator(self):
 
     @pytest.mark.asyncio
     async def test_close_async_delegates_to_chat_generator(self):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         chat_generator.close_async = AsyncMock()
         agent = Agent(chat_generator=chat_generator, tools=[])
         await agent.close_async()
@@ -2471,7 +2399,7 @@ async def test_close_async_delegates_to_chat_generator(self):
 
     @pytest.mark.asyncio
     async def test_close_async_falls_back_to_sync_close(self):
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         chat_generator.close = MagicMock()
         agent = Agent(chat_generator=chat_generator, tools=[])
         await agent.close_async()
@@ -2503,7 +2431,7 @@ class Planner:
             def run(self) -> dict:
                 return {"messages": [ChatMessage.from_assistant("?")], "last_role": "assistant"}
 
-        chat_generator = MockChatGenerator()
+        chat_generator = MockChatGenerator("Hello")
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
         chat_generator.run = MagicMock(return_value={"replies": [ChatMessage.from_assistant("x")]})
 
diff --git a/test/components/agents/test_agent_hitl.py b/test/components/agents/test_agent_hitl.py
index e9dcaac1e69..ad20e47f533 100644
--- a/test/components/agents/test_agent_hitl.py
+++ b/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"}
 
diff --git a/test/components/agents/test_agent_hooks.py b/test/components/agents/test_agent_hooks.py
index 314be7c2800..41cd07b1474 100644
--- a/test/components/agents/test_agent_hooks.py
+++ b/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
diff --git a/test/components/extractors/image/test_llm_document_content_extractor.py b/test/components/extractors/image/test_llm_document_content_extractor.py
index ecad87e28a7..8edaaad7858 100644
--- a/test/components/extractors/image/test_llm_document_content_extractor.py
+++ b/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)
diff --git a/test/components/query/test_query_expander.py b/test/components/query/test_query_expander.py
index 3bd7ded7865..224988aaca1 100644
--- a/test/components/query/test_query_expander.py
+++ b/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
diff --git a/test/components/rankers/test_llm_ranker.py b/test/components/rankers/test_llm_ranker.py
index 077999133fe..683e1211474 100644
--- a/test/components/rankers/test_llm_ranker.py
+++ b/test/components/rankers/test_llm_ranker.py
@@ -9,6 +9,7 @@
 from jinja2 import TemplateSyntaxError
 
 from haystack import Document
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.components.generators.chat.openai import OpenAIChatGenerator
 from haystack.components.rankers.llm_ranker import DEFAULT_PROMPT_TEMPLATE, LLMRanker
 from haystack.dataclasses import ChatMessage
@@ -109,84 +110,78 @@ def test_run_whitespace_query_returns_fallback(mock_chat_generator):
     mock_chat_generator.run.assert_not_called()
 
 
-def test_run_successful_ranking(mock_chat_generator):
+def test_run_successful_ranking():
     documents = [
         Document(id="1", content="first"),
         Document(id="2", content="second"),
         Document(id="3", content="third"),
     ]
-    mock_chat_generator.run.return_value = {
-        "replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 1}, {"index": 3}]}')]
-    }
-    ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=2)
+    chat_generator = MockChatGenerator('{"documents": [{"index": 2}, {"index": 1}, {"index": 3}]}')
+    ranker = LLMRanker(chat_generator=chat_generator, top_k=2)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert [document.id for document in result["documents"]] == ["2", "1"]
 
 
-def test_run_returns_only_documents_listed_by_the_llm(mock_chat_generator):
+def test_run_returns_only_documents_listed_by_the_llm():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}]}')]}
-    ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=2)
+    chat_generator = MockChatGenerator('{"documents": [{"index": 2}]}')
+    ranker = LLMRanker(chat_generator=chat_generator, top_k=2)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert [document.id for document in result["documents"]] == ["2"]
 
 
-def test_run_runtime_top_k_overrides_instance_top_k(mock_chat_generator):
+def test_run_runtime_top_k_overrides_instance_top_k():
     documents = [
         Document(id="doc_1", content="first"),
         Document(id="doc_2", content="second"),
         Document(id="doc_3", content="third"),
     ]
-    mock_chat_generator.run.return_value = {
-        "replies": [ChatMessage.from_assistant('{"documents": [{"index": 3}, {"index": 2}, {"index": 1}]}')]
-    }
-    ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=3)
+    chat_generator = MockChatGenerator('{"documents": [{"index": 3}, {"index": 2}, {"index": 1}]}')
+    ranker = LLMRanker(chat_generator=chat_generator, top_k=3)
 
     result = ranker.run(query="test query", documents=documents, top_k=1)
 
     assert [document.id for document in result["documents"]] == ["doc_3"]
 
 
-def test_run_ignores_out_of_range_indices(mock_chat_generator):
+def test_run_ignores_out_of_range_indices():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {
-        "replies": [ChatMessage.from_assistant('{"documents": [{"index": 99}, {"index": 2}, {"index": 1}]}')]
-    }
-    ranker = LLMRanker(chat_generator=mock_chat_generator)
+    chat_generator = MockChatGenerator('{"documents": [{"index": 99}, {"index": 2}, {"index": 1}]}')
+    ranker = LLMRanker(chat_generator=chat_generator)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert [document.id for document in result["documents"]] == ["2", "1"]
 
 
-def test_run_empty_ranking_result_returns_empty_documents(mock_chat_generator):
+def test_run_empty_ranking_result_returns_empty_documents():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": []}')]}
-    ranker = LLMRanker(chat_generator=mock_chat_generator)
+    chat_generator = MockChatGenerator('{"documents": []}')
+    ranker = LLMRanker(chat_generator=chat_generator)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert result == {"documents": []}
 
 
-def test_run_invalid_json_falls_back(mock_chat_generator):
+def test_run_invalid_json_falls_back():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("not-json")]}
-    ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1, raise_on_failure=False)
+    chat_generator = MockChatGenerator("not-json")
+    ranker = LLMRanker(chat_generator=chat_generator, top_k=1, raise_on_failure=False)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert result == {"documents": documents}
 
 
-def test_run_invalid_json_raises(mock_chat_generator):
+def test_run_invalid_json_raises():
     documents = [Document(id="1", content="first")]
-    mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("not-json")]}
-    ranker = LLMRanker(chat_generator=mock_chat_generator, raise_on_failure=True)
+    chat_generator = MockChatGenerator("not-json")
+    ranker = LLMRanker(chat_generator=chat_generator, raise_on_failure=True)
 
     with pytest.raises(ValueError):
         ranker.run(query="test query", documents=documents)
@@ -221,72 +216,64 @@ def test_run_no_replies_falls_back(mock_chat_generator):
     assert result == {"documents": documents}
 
 
-def test_run_reply_without_text_falls_back(mock_chat_generator):
+def test_run_reply_without_text_falls_back():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant(tool_calls=[])]}
-    ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
+    chat_generator = MockChatGenerator(ChatMessage.from_assistant(tool_calls=[]))
+    ranker = LLMRanker(chat_generator=chat_generator, top_k=1)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert result == {"documents": documents}
 
 
-def test_run_no_valid_document_indices_falls_back(mock_chat_generator):
+def test_run_no_valid_document_indices_falls_back():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {
-        "replies": [ChatMessage.from_assistant('{"documents": [{"index": 0}, {"index": 3}]}')]
-    }
-    ranker = LLMRanker(chat_generator=mock_chat_generator, top_k=1)
+    chat_generator = MockChatGenerator('{"documents": [{"index": 0}, {"index": 3}]}')
+    ranker = LLMRanker(chat_generator=chat_generator, top_k=1)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert result == {"documents": documents}
 
 
-def test_run_deduplicates_documents_before_ranking(mock_chat_generator):
+def test_run_deduplicates_documents_before_ranking():
     documents = [
         Document(id="duplicate", content="keep me", score=0.9),
         Document(id="duplicate", content="drop me", score=0.1),
         Document(id="unique", content="unique", score=0.2),
     ]
-    mock_chat_generator.run.return_value = {
-        "replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 1}]}')]
-    }
-    ranker = LLMRanker(chat_generator=mock_chat_generator)
+    chat_generator = MockChatGenerator('{"documents": [{"index": 2}, {"index": 1}]}')
+    ranker = LLMRanker(chat_generator=chat_generator)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert [document.content for document in result["documents"]] == ["unique", "keep me"]
 
 
-def test_run_preserves_duplicate_indices(mock_chat_generator):
+def test_run_preserves_duplicate_indices():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {
-        "replies": [ChatMessage.from_assistant('{"documents": [{"index": 2}, {"index": 2}, {"index": 1}]}')]
-    }
-    ranker = LLMRanker(chat_generator=mock_chat_generator)
+    chat_generator = MockChatGenerator('{"documents": [{"index": 2}, {"index": 2}, {"index": 1}]}')
+    ranker = LLMRanker(chat_generator=chat_generator)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert [document.id for document in result["documents"]] == ["2", "2", "1"]
 
 
-def test_run_numeric_string_index_is_accepted(mock_chat_generator):
+def test_run_numeric_string_index_is_accepted():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant('{"documents": [{"index": "2"}]}')]}
-    ranker = LLMRanker(chat_generator=mock_chat_generator)
+    chat_generator = MockChatGenerator('{"documents": [{"index": "2"}]}')
+    ranker = LLMRanker(chat_generator=chat_generator)
 
     result = ranker.run(query="test query", documents=documents)
 
     assert result == {"documents": [documents[1]]}
 
 
-def test_run_invalid_index_type_falls_back(mock_chat_generator):
+def test_run_invalid_index_type_falls_back():
     documents = [Document(id="1", content="first"), Document(id="2", content="second")]
-    mock_chat_generator.run.return_value = {
-        "replies": [ChatMessage.from_assistant('{"documents": [{"index": "invalid"}]}')]
-    }
-    ranker = LLMRanker(chat_generator=mock_chat_generator)
+    chat_generator = MockChatGenerator('{"documents": [{"index": "invalid"}]}')
+    ranker = LLMRanker(chat_generator=chat_generator)
 
     result = ranker.run(query="test query", documents=documents)
 
diff --git a/test/core/pipeline/breakpoints/test_pipeline_breakpoints_answer_joiner.py b/test/core/pipeline/breakpoints/test_pipeline_breakpoints_answer_joiner.py
index 22c5ac804d9..9ac973e1241 100644
--- a/test/core/pipeline/breakpoints/test_pipeline_breakpoints_answer_joiner.py
+++ b/test/core/pipeline/breakpoints/test_pipeline_breakpoints_answer_joiner.py
@@ -4,8 +4,8 @@
 
 import pytest
 
-from haystack import component
 from haystack.components.builders.answer_builder import AnswerBuilder
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.components.joiners import AnswerJoiner
 from haystack.core.errors import BreakpointException
 from haystack.core.pipeline.pipeline import Pipeline
@@ -13,24 +13,13 @@
 from haystack.dataclasses.breakpoints import Breakpoint
 
 
-@component
-class FakeChatGenerator:
-    def __init__(self, content: str, model_name: str):
-        self.content = content
-        self.model_name = model_name
-
-    @component.output_types(replies=list[ChatMessage])
-    def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
-        return {"replies": [ChatMessage.from_assistant(self.content)]}
-
-
 class TestPipelineBreakpoints:
     @pytest.fixture
     def answer_join_pipeline(self):
         """Creates a pipeline with fake components."""
         pipeline = Pipeline()
-        pipeline.add_component("gpt-4o", FakeChatGenerator("GPT-4 response", "gpt-4o"))
-        pipeline.add_component("gpt-3", FakeChatGenerator("GPT-3 response", "gpt-3.5-turbo"))
+        pipeline.add_component("gpt-4o", MockChatGenerator("GPT-4 response", model="gpt-4o"))
+        pipeline.add_component("gpt-3", MockChatGenerator("GPT-3 response", model="gpt-3.5-turbo"))
         pipeline.add_component("answer_builder_a", AnswerBuilder())
         pipeline.add_component("answer_builder_b", AnswerBuilder())
         pipeline.add_component("answer_joiner", AnswerJoiner())
diff --git a/test/core/pipeline/breakpoints/test_pipeline_breakpoints_branch_joiner.py b/test/core/pipeline/breakpoints/test_pipeline_breakpoints_branch_joiner.py
index 03b5e602f83..68c1ae43b32 100644
--- a/test/core/pipeline/breakpoints/test_pipeline_breakpoints_branch_joiner.py
+++ b/test/core/pipeline/breakpoints/test_pipeline_breakpoints_branch_joiner.py
@@ -3,12 +3,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import json
-from typing import Any
 
 import pytest
 
-from haystack import component
 from haystack.components.converters import OutputAdapter
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.components.joiners import BranchJoiner
 from haystack.components.validators import JsonSchemaValidator
 from haystack.core.errors import BreakpointException
@@ -17,18 +16,6 @@
 from haystack.dataclasses.breakpoints import Breakpoint
 
 
-@component
-class FakeChatGenerator:
-    def __init__(self, content: str):
-        self.content = content
-
-    @component.output_types(replies=list[ChatMessage])
-    def run(
-        self, messages: list[ChatMessage], generation_kwargs: dict | None = None, **kwargs: Any
-    ) -> dict[str, list[ChatMessage]]:
-        return {"replies": [ChatMessage.from_assistant(self.content)]}
-
-
 class TestPipelineBreakpoints:
     @pytest.fixture
     def branch_joiner_pipeline(self):
@@ -46,7 +33,7 @@ def branch_joiner_pipeline(self):
 
         pipe = Pipeline()
         pipe.add_component("joiner", BranchJoiner(list[ChatMessage]))
-        pipe.add_component("fc_llm", FakeChatGenerator(content))
+        pipe.add_component("fc_llm", MockChatGenerator(content))
         pipe.add_component("validator", JsonSchemaValidator(json_schema=person_schema))
         pipe.add_component("adapter", OutputAdapter("{{chat_message}}", list[ChatMessage], unsafe=True))
 
diff --git a/test/core/pipeline/breakpoints/test_pipeline_breakpoints_list_joiner.py b/test/core/pipeline/breakpoints/test_pipeline_breakpoints_list_joiner.py
index 68a488cff08..210860a9099 100644
--- a/test/core/pipeline/breakpoints/test_pipeline_breakpoints_list_joiner.py
+++ b/test/core/pipeline/breakpoints/test_pipeline_breakpoints_list_joiner.py
@@ -2,28 +2,17 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any
-
 import pytest
 
-from haystack import Pipeline, component
+from haystack import Pipeline
 from haystack.components.builders import ChatPromptBuilder
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.components.joiners import ListJoiner
 from haystack.core.errors import BreakpointException
 from haystack.dataclasses import ChatMessage
 from haystack.dataclasses.breakpoints import Breakpoint
 
 
-@component
-class FakeChatGenerator:
-    def __init__(self, response: str):
-        self.response = response
-
-    @component.output_types(replies=list[ChatMessage])
-    def run(self, messages: list[ChatMessage], **kwargs: Any) -> dict[str, list[ChatMessage]]:
-        return {"replies": [ChatMessage.from_assistant(self.response)]}
-
-
 class TestPipelineBreakpoints:
     @pytest.fixture
     def list_joiner_pipeline(self):
@@ -39,11 +28,11 @@ def list_joiner_pipeline(self):
 
         pipe = Pipeline()
         pipe.add_component("prompt_builder", ChatPromptBuilder(template=user_message, required_variables=None))
-        pipe.add_component("llm", FakeChatGenerator("Nuclear physics is the study of atomic nuclei."))
+        pipe.add_component("llm", MockChatGenerator("Nuclear physics is the study of atomic nuclei."))
         pipe.add_component(
             "feedback_prompt_builder", ChatPromptBuilder(template=feedback_message, required_variables=None)
         )
-        pipe.add_component("feedback_llm", FakeChatGenerator("Score: 8/10. Concise and accurate."))
+        pipe.add_component("feedback_llm", MockChatGenerator("Score: 8/10. Concise and accurate."))
         pipe.add_component("list_joiner", ListJoiner(list[ChatMessage]))
 
         pipe.connect("prompt_builder.prompt", "llm.messages")
diff --git a/test/core/pipeline/breakpoints/test_pipeline_breakpoints_loops.py b/test/core/pipeline/breakpoints/test_pipeline_breakpoints_loops.py
index d1832284e69..99b9f5d3091 100644
--- a/test/core/pipeline/breakpoints/test_pipeline_breakpoints_loops.py
+++ b/test/core/pipeline/breakpoints/test_pipeline_breakpoints_loops.py
@@ -10,6 +10,7 @@
 
 from haystack import component
 from haystack.components.builders import ChatPromptBuilder
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.core.errors import BreakpointException
 from haystack.core.pipeline.pipeline import Pipeline
 from haystack.dataclasses import ChatMessage
@@ -34,16 +35,6 @@ def run(self, replies: list[ChatMessage]) -> dict[str, list[ChatMessage] | str]:
             return {"invalid_replies": replies, "error_message": str(e)}
 
 
-@component
-class FakeChatGenerator:
-    def __init__(self, response: str):
-        self.response = response
-
-    @component.output_types(replies=list[ChatMessage])
-    def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
-        return {"replies": [ChatMessage.from_assistant(self.response)]}
-
-
 class City(BaseModel):
     name: str
     country: str
@@ -95,7 +86,7 @@ def validation_loop_pipeline(self):
             instance=ChatPromptBuilder(template=prompt_template, required_variables=["passage", "schema"]),
             name="prompt_builder",
         )
-        pipeline.add_component(instance=FakeChatGenerator(response=response_json), name="llm")
+        pipeline.add_component(instance=MockChatGenerator(response_json), name="llm")
         pipeline.add_component(instance=OutputValidator(pydantic_model=CitiesData), name="output_validator")
 
         pipeline.connect("prompt_builder.prompt", "llm.messages")
diff --git a/test/core/pipeline/test_pipeline.py b/test/core/pipeline/test_pipeline.py
index 93328513a5e..1c4da6962eb 100644
--- a/test/core/pipeline/test_pipeline.py
+++ b/test/core/pipeline/test_pipeline.py
@@ -8,18 +8,12 @@
 import pytest
 
 from haystack.components.agents import Agent
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.components.joiners import BranchJoiner
 from haystack.core.component import component
 from haystack.core.errors import PipelineConnectError, PipelineRuntimeError
 from haystack.core.pipeline import Pipeline
-from haystack.dataclasses import ChatMessage, Document
-
-
-@component
-class MockChatGenerator:
-    @component.output_types(replies=list[ChatMessage])
-    def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
-        return {"replies": [ChatMessage.from_assistant("Hello, world!")]}
+from haystack.dataclasses import ChatMessage, ChatRole, Document
 
 
 @component
@@ -400,16 +394,16 @@ def run(self) -> dict[str, list[ChatMessage]]:
         p = Pipeline()
         p.add_component("message_producer", MessageProducer())
         p.add_component("message_producer2", MessageProducer())
-        p.add_component("agent", Agent(chat_generator=MockChatGenerator()))
+        p.add_component("agent", Agent(chat_generator=MockChatGenerator("Hello, world!")))
         p.connect("message_producer", "agent.messages")
         p.connect("message_producer2", "agent.messages")
 
         result = p.run({})
-        assert result["agent"]["messages"] == [
-            ChatMessage.from_user("Hello, world!"),
-            ChatMessage.from_user("Hello, world!"),
-            ChatMessage.from_assistant("Hello, world!"),
-        ]
+        messages = result["agent"]["messages"]
+
+        assert messages[:2] == [ChatMessage.from_user("Hello, world!"), ChatMessage.from_user("Hello, world!")]
+        assert messages[2].role == ChatRole.ASSISTANT
+        assert messages[2].text == "Hello, world!"
 
     def test_run_auto_variadic_str_to_list_str(self):
         """Two str producers connected to a list[str] input are auto-joined and flattened at runtime."""
diff --git a/test/hooks/tool_result_offloading/test_hooks.py b/test/hooks/tool_result_offloading/test_hooks.py
index 4229183e8c0..2b302d1dba1 100644
--- a/test/hooks/tool_result_offloading/test_hooks.py
+++ b/test/hooks/tool_result_offloading/test_hooks.py
@@ -3,14 +3,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from pathlib import Path
-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.state import State
+from haystack.components.generators.chat import MockChatGenerator
 from haystack.dataclasses import ChatMessage, ImageContent, TextContent, ToolCall
 from haystack.hooks.tool_result_offloading import (
     RESULT_STORE_CONTEXT_KEY,
@@ -20,20 +20,7 @@
     OffloadOverChars,
     ToolResultOffloadHook,
 )
-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("done")]}
-
-    @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("done")]}
+from haystack.tools import tool
 
 
 @tool
@@ -237,7 +224,7 @@ def test_offloads_tool_result_seen_by_next_llm_call(self, tmp_path):
         hook = ToolResultOffloadHook(
             store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()}
         )
-        agent = Agent(chat_generator=MockChatGenerator(), tools=[big_tool], hooks={"after_tool": [hook]})
+        agent = Agent(chat_generator=MockChatGenerator("done"), tools=[big_tool], hooks={"after_tool": [hook]})
         agent.warm_up()
         agent.chat_generator.run = MagicMock(
             side_effect=[
@@ -258,7 +245,7 @@ async def test_offloads_tool_result_async(self, tmp_path):
         hook = ToolResultOffloadHook(
             store=FileSystemToolResultStore(root=tmp_path), offload_strategies={"*": AlwaysOffload()}
         )
-        agent = Agent(chat_generator=MockChatGenerator(), tools=[big_tool], hooks={"after_tool": [hook]})
+        agent = Agent(chat_generator=MockChatGenerator("done"), tools=[big_tool], hooks={"after_tool": [hook]})
         agent.warm_up()
         agent.chat_generator.run_async = AsyncMock(
             side_effect=[
diff --git a/test/tools/test_component_tool.py b/test/tools/test_component_tool.py
index dd0c42bc5ba..ea4199476b1 100644
--- a/test/tools/test_component_tool.py
+++ b/test/tools/test_component_tool.py
@@ -15,10 +15,10 @@
 from haystack import Pipeline, SuperComponent, component
 from haystack.components.agents import Agent, State
 from haystack.components.builders import PromptBuilder
-from haystack.components.generators.chat import OpenAIChatGenerator
+from haystack.components.generators.chat import MockChatGenerator, OpenAIChatGenerator
 from haystack.core.pipeline.utils import _deepcopy_with_exceptions
 from haystack.dataclasses import ChatMessage, ChatRole, Document
-from haystack.tools import ComponentTool, ToolsType
+from haystack.tools import ComponentTool
 from test.tools.test_parameters_schema_utils import BYTE_STREAM_SCHEMA, DOCUMENT_SCHEMA, SPARSE_EMBEDDING_SCHEMA
 
 # Component and Model Definitions
@@ -148,22 +148,6 @@ def run(self, documents: list[Document], top_k: int = 5) -> dict[str, str]:
         return {"concatenated": "\n".join(doc.content for doc in documents[:top_k] if doc.content)}
 
 
-@component
-class FakeChatGenerator:
-    def __init__(self, messages: list[ChatMessage]):
-        self.messages = messages
-
-    @component.output_types(replies=list[ChatMessage])
-    def run(
-        self,
-        messages: list[ChatMessage],
-        generation_kwargs: dict[str, Any] | None = None,
-        *,
-        tools: ToolsType | None = None,
-    ) -> dict[str, list[ChatMessage]]:
-        return {"replies": self.messages}
-
-
 def output_handler(old, new):
     """
     Output handler to test serialization.
@@ -555,7 +539,7 @@ def run(self, query: str, state: State | None = None) -> dict:
 
     def test_component_invoker_with_agent(self):
         """Tests that Agent as a ComponentTool can be invoked when calling it with a list of dicts"""
-        agent = Agent(chat_generator=FakeChatGenerator(messages=[ChatMessage.from_assistant("Answer")]))
+        agent = Agent(chat_generator=MockChatGenerator("Answer"))
         tool = ComponentTool(
             component=agent,
             name="agent_tool",
@@ -563,7 +547,7 @@ def test_component_invoker_with_agent(self):
             outputs_to_string={"source": "last_message"},
         )
         result = tool.invoke(messages=[{"role": "user", "content": [{"text": "A 4-day trip in the south of France"}]}])
-        assert result["last_message"] == ChatMessage.from_assistant("Answer")
+        assert result["last_message"].text == "Answer"
 
     def test_convert_param_union_with_list_arm(self):
         @component
diff --git a/test/tools/test_searchable_toolset.py b/test/tools/test_searchable_toolset.py
index 1bf0d5155f8..c83854c7a15 100644
--- a/test/tools/test_searchable_toolset.py
+++ b/test/tools/test_searchable_toolset.py
@@ -11,7 +11,7 @@
 
 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.tools import SearchableToolset, Tool, Toolset, flatten_tools_or_toolsets
 from haystack.tools.from_function import create_tool_from_function
@@ -923,32 +923,18 @@ def test_discovered_tool_call_counts_added_lazily(self, large_catalog):
         """tool_call_counts seeds only the search tool up front; a discovered+called tool is added lazily."""
         toolset = SearchableToolset(catalog=large_catalog, search_threshold=3, top_k=5)
 
-        @component
-        class SearchThenWeatherGenerator:
-            step = 0
-
-            @component.output_types(replies=list[ChatMessage])
-            def run(self, messages, tools=None, **kwargs):
-                self.step += 1
-                if self.step == 1:
-                    return {
-                        "replies": [
-                            ChatMessage.from_assistant(
-                                tool_calls=[ToolCall(tool_name="search_tools", arguments={"tool_keywords": "weather"})]
-                            )
-                        ]
-                    }
-                if self.step == 2:
-                    return {
-                        "replies": [
-                            ChatMessage.from_assistant(
-                                tool_calls=[ToolCall(tool_name="get_weather", arguments={"city": "Berlin"})]
-                            )
-                        ]
-                    }
-                return {"replies": [ChatMessage.from_assistant("done")]}
-
-        agent = Agent(chat_generator=SearchThenWeatherGenerator(), tools=toolset, max_agent_steps=6)
+        chat_generator = MockChatGenerator(
+            [
+                ChatMessage.from_assistant(
+                    tool_calls=[ToolCall(tool_name="search_tools", arguments={"tool_keywords": "weather"})]
+                ),
+                ChatMessage.from_assistant(
+                    tool_calls=[ToolCall(tool_name="get_weather", arguments={"city": "Berlin"})]
+                ),
+                "done",
+            ]
+        )
+        agent = Agent(chat_generator=chat_generator, tools=toolset, max_agent_steps=6)
         result = agent.run(messages=[ChatMessage.from_user("What's the weather in Berlin?")])
 
         counts = result["tool_call_counts"]