63 lines
1.8 KiB
Python
63 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from unittest.mock import MagicMock
|
|
|
|
from langchain_core.messages import ToolMessage
|
|
|
|
from deerflow.agents.middlewares.sensitive_output_redaction_middleware import SensitiveOutputRedactionMiddleware
|
|
|
|
|
|
def _request(name: str = "bash", args: dict | None = None):
|
|
req = MagicMock()
|
|
req.tool_call = {"name": name, "args": args or {}, "id": "call_1"}
|
|
return req
|
|
|
|
|
|
def test_redacts_key_value_and_bearer():
|
|
mw = SensitiveOutputRedactionMiddleware()
|
|
req = _request()
|
|
handler = MagicMock(
|
|
return_value=ToolMessage(
|
|
content="OPENAI_API_KEY=sk-abc123456789\nAuthorization: Bearer abcdefghijklmnop",
|
|
tool_call_id="call_1",
|
|
name="bash",
|
|
status="success",
|
|
)
|
|
)
|
|
result = mw.wrap_tool_call(req, handler)
|
|
assert "OPENAI_API_KEY=[REDACTED]" in result.content
|
|
assert "Bearer [REDACTED]" in result.content
|
|
|
|
|
|
def test_redacts_private_key_block():
|
|
mw = SensitiveOutputRedactionMiddleware()
|
|
req = _request()
|
|
handler = MagicMock(
|
|
return_value=ToolMessage(
|
|
content="-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----",
|
|
tool_call_id="call_1",
|
|
name="bash",
|
|
status="success",
|
|
)
|
|
)
|
|
result = mw.wrap_tool_call(req, handler)
|
|
assert result.content == "[REDACTED]"
|
|
|
|
|
|
def test_async_path_redacts_jwt():
|
|
mw = SensitiveOutputRedactionMiddleware()
|
|
req = _request()
|
|
|
|
async def handler(_):
|
|
return ToolMessage(
|
|
content="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.abcde12345.zyxwv98765",
|
|
tool_call_id="call_1",
|
|
name="bash",
|
|
status="success",
|
|
)
|
|
|
|
result = asyncio.run(mw.awrap_tool_call(req, handler))
|
|
assert result.content == "[REDACTED]"
|
|
|