@@ -41,7 +41,7 @@ def factorial_failing(n: int) -> dict:
4141
4242
4343@component
44- class TestChatGenerator :
44+ class MockChatGenerator :
4545 def __init__ (self , fail_on_call : bool ):
4646 self .fail_on_call = fail_on_call
4747
@@ -85,7 +85,7 @@ def test_pipeline_with_chat_generator_crash():
8585 """Test pipeline crash handling when chat generator fails."""
8686 pipe = build_pipeline (
8787 agent = Agent (
88- chat_generator = TestChatGenerator (True ), tools = [calculator_tool ], state_schema = {"calc_result" : {"type" : int }}
88+ chat_generator = MockChatGenerator (True ), tools = [calculator_tool ], state_schema = {"calc_result" : {"type" : int }}
8989 )
9090 )
9191
@@ -96,7 +96,7 @@ def test_pipeline_with_chat_generator_crash():
9696
9797 assert "Error in chat generator component" in str (exception_info .value )
9898 assert exception_info .value .component_name == "chat_generator"
99- assert exception_info .value .component_type == TestChatGenerator
99+ assert exception_info .value .component_type == MockChatGenerator
100100 assert "math_agent_chat_generator" in exception_info .value .pipeline_snapshot_file_path
101101
102102 pipeline_snapshot = exception_info .value .pipeline_snapshot
@@ -131,7 +131,7 @@ def test_pipeline_with_tool_call_crash():
131131 """Test pipeline crash handling when a tool call fails."""
132132 pipe = build_pipeline (
133133 agent = Agent (
134- chat_generator = TestChatGenerator (False ),
134+ chat_generator = MockChatGenerator (False ),
135135 tools = [calculator_tool , failing_factorial_tool ],
136136 state_schema = {"calc_result" : {"type" : int }, "factorial_result" : {"type" : int }},
137137 raise_on_tool_invocation_failure = True ,
0 commit comments