Home / Class/ TestPIIMiddlewareIntegration Class — langchain Architecture

TestPIIMiddlewareIntegration Class — langchain Architecture

Architecture documentation for the TestPIIMiddlewareIntegration class in test_pii.py from the langchain codebase.

Entity Profile

Dependency Diagram

graph TD
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a["TestPIIMiddlewareIntegration"]
  511a8f91_080f_9750_8245_6c31ee0b22eb["ToolMessage"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|extends| 511a8f91_080f_9750_8245_6c31ee0b22eb
  2664a7f2_5314_35a9_a86d_2fce142dd0cf["PIIDetectionError"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|extends| 2664a7f2_5314_35a9_a86d_2fce142dd0cf
  6dcffc03_fae9_01f8_ed9b_d30d8a73dd13["test_pii.py"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|defined in| 6dcffc03_fae9_01f8_ed9b_d30d8a73dd13
  20ab26af_f02d_055e_5ebd_45c47423fc7a["test_apply_to_input_only()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| 20ab26af_f02d_055e_5ebd_45c47423fc7a
  ab248ac4_fcae_aca0_054b_7c971a1219a2["test_apply_to_output_only()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| ab248ac4_fcae_aca0_054b_7c971a1219a2
  996e9157_5e22_ce83_4ef7_c9f4cc02a1f1["test_apply_to_both()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| 996e9157_5e22_ce83_4ef7_c9f4cc02a1f1
  0adf1db3_cbc6_c082_9e82_609ea365b291["test_no_pii_returns_none()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| 0adf1db3_cbc6_c082_9e82_609ea365b291
  9035ca96_dbe2_8e56_e1ca_e7151ce1798c["test_empty_messages()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| 9035ca96_dbe2_8e56_e1ca_e7151ce1798c
  e26537eb_1a70_dfc2_feae_9d059b92f641["test_apply_to_tool_results()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| e26537eb_1a70_dfc2_feae_9d059b92f641
  472fa60a_3e95_c195_f4fa_551561690f94["test_apply_to_tool_results_mask_strategy()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| 472fa60a_3e95_c195_f4fa_551561690f94
  4cbd77cc_a2ca_e28c_153f_3f893dc106c9["test_apply_to_tool_results_block_strategy()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| 4cbd77cc_a2ca_e28c_153f_3f893dc106c9
  5203b000_daf4_f837_b642_40dff2749f89["test_with_agent()"]
  74acad26_e3cd_8385_fcb7_9bf787fa3f4a -->|method| 5203b000_daf4_f837_b642_40dff2749f89

Relationship Graph

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_pii.py lines 354–512

class TestPIIMiddlewareIntegration:
    """Test PIIMiddleware integration with agent."""

    def test_apply_to_input_only(self) -> None:
        """Test that middleware only processes input when configured."""
        middleware = PIIMiddleware(
            "email", strategy="redact", apply_to_input=True, apply_to_output=False
        )

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

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

    def test_apply_to_output_only(self) -> None:
        """Test that middleware only processes output when configured."""
        middleware = PIIMiddleware(
            "email", strategy="redact", apply_to_input=False, apply_to_output=True
        )

        # Should not process HumanMessage
        state = AgentState[Any](messages=[HumanMessage("Email: test@example.com")])
        result = middleware.before_model(state, Runtime())
        assert result is 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
        assert "[REDACTED_EMAIL]" in result["messages"][0].content

    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

    def test_no_pii_returns_none(self) -> None:
        """Test that middleware returns None when no PII detected."""
        middleware = PIIMiddleware("email", strategy="redact")
        state = AgentState[Any](messages=[HumanMessage("No PII here")])

        result = middleware.before_model(state, Runtime())
        assert result is None

    def test_empty_messages(self) -> None:
        """Test that middleware handles empty messages gracefully."""
        middleware = PIIMiddleware("email", strategy="redact")
        state = AgentState[Any](messages=[])

        result = middleware.before_model(state, Runtime())
        assert result is None

    def test_apply_to_tool_results(self) -> None:
        """Test that middleware processes tool results when enabled."""
        middleware = PIIMiddleware(
            "email", strategy="redact", apply_to_input=False, apply_to_tool_results=True
        )

        # Simulate a conversation with tool call and result containing PII
        state = AgentState[Any](
            messages=[
                HumanMessage("Search for John"),
                AIMessage(
                    content="",

Frequently Asked Questions

What is the TestPIIMiddlewareIntegration class?
TestPIIMiddlewareIntegration is a class in the langchain codebase, defined in libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_pii.py.
Where is TestPIIMiddlewareIntegration defined?
TestPIIMiddlewareIntegration is defined in libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_pii.py at line 354.
What does TestPIIMiddlewareIntegration extend?
TestPIIMiddlewareIntegration extends ToolMessage, PIIDetectionError.

Analyze Your Own Codebase

Get architecture documentation, dependency graphs, and domain analysis for your codebase in minutes.

Try Supermodel Free