diff --git a/tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py b/tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py index f2f7e35b05..43474ad353 100644 --- a/tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py +++ b/tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py @@ -171,7 +171,7 @@ def test_on_model_error_callback_with_plugin(mock_plugin): def test_on_model_error_callback_fallback_to_runner(mock_plugin): - """Tests that the model error is not handled and falls back to raise from runner.""" + """Tests that the model error is not handled and surfaces from the runner.""" mock_model = testing_utils.MockModel.create(error=mock_error, responses=[]) mock_plugin.enable_on_model_error_callback = False agent = Agent( @@ -179,10 +179,12 @@ def test_on_model_error_callback_fallback_to_runner(mock_plugin): model=mock_model, ) - try: - testing_utils.InMemoryRunner(agent, plugins=[mock_plugin]) - except Exception as e: - assert e == mock_error + runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin]) + + events = runner.run('test') + error_events = [e for e in events if e.error_code] + assert len(error_events) == 1 + assert error_events[0].error_code == 'ClientError' if __name__ == '__main__': diff --git a/tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py b/tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py index 3c39e2844b..ff18c38aba 100644 --- a/tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py +++ b/tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py @@ -180,10 +180,9 @@ async def test_async_on_tool_error_fallback_to_runner( ): mock_plugin.enable_on_tool_error_callback = False - try: + with pytest.raises(ClientError) as exc_info: await invoke_tool_with_plugin(mock_error_tool, mock_plugin) - except Exception as e: - assert e == mock_error + assert exc_info.value == mock_error async def invoke_tool_with_plugin_live( @@ -258,10 +257,9 @@ async def test_live_on_tool_error_fallback_to_runner( ): mock_plugin.enable_on_tool_error_callback = False - try: + with pytest.raises(ClientError) as exc_info: await invoke_tool_with_plugin_live(mock_error_tool, mock_plugin) - except Exception as e: - assert e == mock_error + assert exc_info.value == mock_error @pytest.mark.asyncio