ai-agent-book/chapter3/log-sanitization/tests/test_sanitize_conversation.py

#!/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"