ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/log-sanitization/tests/test_sanitize_conversation.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+"""Regression tests for sanitize_conversation() in agent.py.
+
+Bug: detect_pii() catches any backend exception and returns ([], {}), but
+sanitize_conversation then subscripted the empty metrics dict
+(perf_metrics['input_tokens']) -> KeyError that killed the whole batch.
+Fixed with .get(..., 0) defaults so a dead backend degrades gracefully.
+"""
+
+import pytest
+
+from agent import LogSanitizationAgent
+from metrics import MetricsCollector
+
+
+def _make_agent(client, tmp_path):
+    """Build an agent without __init__ (which requires a live Ollama)."""
+    ag = LogSanitizationAgent.__new__(LogSanitizationAgent)
+    ag.model = "qwen3:0.6b"
+    ag.backend = "ollama"
+    ag.metrics_collector = MetricsCollector(tmp_path)
+    ag.client = client
+    return ag
+
+
+CONV = {
+    "conversation_id": "demo_001",
+    "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
+}
+
+
+class _DeadClient:
+    def chat(self, **kwargs):
+        raise ConnectionError("[test] Ollama server is not running")
+
+
+class _FakeClient:
+    """Mimics ollama.Client.chat(stream=True) chunk shape."""
+
+    def __init__(self, payload: str):
+        self.payload = payload
+
+    def chat(self, **kwargs):
+        return [{"message": {"content": self.payload}}]
+
+
+def test_dead_backend_returns_result_with_zero_metrics(tmp_path):
+    ag = _make_agent(_DeadClient(), tmp_path)
+    result = ag.sanitize_conversation(CONV, "t1")  # must not raise
+    assert result["pii_found"] == []
+    assert result["replacements_made"] == 0
+    assert result["metrics"]["input_tokens"] == 0
+    assert result["metrics"]["pii_items_found"] == 0
+
+
+def test_working_backend_still_detects_pii(tmp_path):
+    ag = _make_agent(_FakeClient('{"pii_values": ["123-45-6789"]}'), tmp_path)
+    result = ag.sanitize_conversation(CONV, "t1")
+    assert result["pii_found"] == ["123-45-6789"]
+    assert result["replacements_made"] == 1
+    assert "[REDACTED]" in result["sanitized_text"]
+    assert result["metrics"]["pii_items_found"] == 1
+
+
+def test_working_backend_with_structured_pii_items(tmp_path):
+    conv = {
+        "conversation_id": "demo_002",
+        "messages": [
+            {
+                "role": "user",
+                "content": "My SSN is 000-00-0000 and my card is 0000-0000-0000-0000.",
+            }
+        ],
+    }
+    payload = (
+        '{"pii_items": ['
+        '{"type": "social_security_number", "value": "000-00-0000"}, '
+        '{"type": "credit_card_number", "value": "0000-0000-0000-0000"}'
+        ']}'
+    )
+    ag = _make_agent(_FakeClient(payload), tmp_path)
+    result = ag.sanitize_conversation(conv, "t2")
+    assert result["pii_found"] == ["000-00-0000", "0000-0000-0000-0000"]
+    assert result["replacements_made"] == 2
+    assert result["sanitized_text"].count("[REDACTED]") == 2
+    assert result["metrics"]["pii_items_found"] == 2
+
+
+def test_structured_pii_items_preserve_original_value(tmp_path):
+    conv = {
+        "conversation_id": "demo_003",
+        "messages": [
+            {
+                "role": "user",
+                "content": "The secret is -abc- and the password is   p4ss  .",
+            }
+        ],
+    }
+    payload = (
+        '{"pii_items": ['
+        '{"type": "secret", "value": "-abc-"}, '
+        '{"type": "password", "value": "  p4ss  "}'
+        ']}'
+    )
+    ag = _make_agent(_FakeClient(payload), tmp_path)
+    result = ag.sanitize_conversation(conv, "t3")
+    assert result["pii_found"] == ["-abc-", "  p4ss  "]
+    assert result["replacements_made"] == 2
+    assert result["sanitized_text"].count("[REDACTED]") == 2
+
+
+def test_pii_items_metric_excludes_rejected_and_malformed_items(tmp_path):
+    conv = {
+        "conversation_id": "demo_004",
+        "messages": [
+            {
+                "role": "user",
+                "content": "My email is alice@example.com.",
+            }
+        ],
+    }
+    payload = (
+        '{"pii_items": ['
+        '{"type": "email", "value": "alice@example.com"}, '
+        '{"type": "ssn", "value": "999-99-9999"}, '
+        '{"type": "unknown", "value": ""}, '
+        '{"type": "broken"}, '
+        '"just a string"'
+        ']}'
+    )
+    ag = _make_agent(_FakeClient(payload), tmp_path)
+    result = ag.sanitize_conversation(conv, "t4")
+    assert result["pii_found"] == ["alice@example.com"]
+    assert result["replacements_made"] == 1
+    assert result["metrics"]["pii_items_found"] == 1
+    items = result["pii_items"]
+    assert len(items) == 1
+    assert items[0]["type"] == "email"
+    assert items[0]["value"] == "alice@example.com"