test_error_handling_with_success_path() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 445–493

    def test_error_handling_with_success_path(self) -> None:
        """Test that error handling middleware works correctly on both success and error paths."""
        call_log = []

        class ErrorRecoveryMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                try:
                    call_log.append("before-yield")
                    result = handler(request)
                    call_log.append("after-yield-success")
                except Exception:
                    call_log.append("caught-error")
                    return AIMessage(content="Recovered from error")
                return result

        # Test 1: Success path
        call_log.clear()
        model1 = GenericFakeChatModel(messages=iter([AIMessage(content="Success")]))
        agent1 = create_agent(model=model1, middleware=[ErrorRecoveryMiddleware()])
        result1 = agent1.invoke({"messages": [HumanMessage("Test")]})

        assert result1["messages"][1].content == "Success"
        assert call_log == ["before-yield", "after-yield-success"]

        # Test 2: Error path
        call_log.clear()

        class AlwaysFailModel(GenericFakeChatModel):
            @override
            def _generate(
                self,
                messages: list[BaseMessage],
                stop: list[str] | None = None,
                run_manager: CallbackManagerForLLMRun | None = None,
                **kwargs: Any,
            ) -> ChatResult:
                msg = "Model error"
                raise ValueError(msg)

        model2 = AlwaysFailModel(messages=iter([]))
        agent2 = create_agent(model=model2, middleware=[ErrorRecoveryMiddleware()])
        result2 = agent2.invoke({"messages": [HumanMessage("Test")]})

        assert result2["messages"][1].content == "Recovered from error"
        assert call_log == ["before-yield", "caught-error"]