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="",
Extends
Source
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