test_apply_to_both() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_pii.py lines 391–405

    def test_apply_to_both(self) -> None:
        """Test that middleware processes both input and output."""
        middleware = PIIMiddleware(
            "email", strategy="redact", apply_to_input=True, apply_to_output=True
        )

        # Should process HumanMessage
        state = AgentState[Any](messages=[HumanMessage("Email: test@example.com")])
        result = middleware.before_model(state, Runtime())
        assert result is not None

        # Should process AIMessage
        state = AgentState[Any](messages=[AIMessage("My email is ai@example.com")])
        result = middleware.after_model(state, Runtime())
        assert result is not None