tests/observability/test_sanitize.py

"""Redaction tests for the observability sanitizer.

The forbidden-content list is copied verbatim from the shared plan preamble's
global constraints: "Never export prompts, full responses, PDF text, chunks,
annotation contents, filesystem paths, URL query strings, base64, credentials,
API headers, or raw document IDs." Each item below has at least one test.
"""

from __future__ import annotations

import pytest

from app.observability.sanitize import (
    REDACTED,
    is_forbidden_key,
    sanitize_attribute_value,
    sanitize_attributes,
    sanitize_exception_text,
)


# --- Forbidden content keys (prompts, responses, chunks, PDF, annotations) ---


@pytest.mark.parametrize(
    "key",
    [
        "prompt",
        "system_prompt",
        "llm_response",
        "response_text",
        "chunk",
        "retrieved_chunks",
        "pdf_text",
        "annotation",
        "annotation_body",
        "content",
        "selection_text",
        "message_body",
    ],
)
def test_forbidden_content_keys_are_recognized(key):
    assert is_forbidden_key(key) is True


@pytest.mark.parametrize(
    "key",
    [
        "raw_candidate_count",
        "usable_candidate_count",
        "requested_top_k",
        "collection_size",
        "duplicate_candidate_ratio",
        "document_diversity",
        "section_diversity",
        "context_input_chars",
        "context_selected_chars",
        "context_truncated",
        "cache.state",
        "pipeline.version",
        "query.category",
        "retrieval_status",
        "empty_result",
        "retrieval_error",
    ],
)
def test_metric_vocabulary_keys_are_not_forbidden(key):
    # These are the real retrieval/experiment measures -- token-splitting must
    # not false-positive on "context" (ends in "text") or "document".
    assert is_forbidden_key(key) is False


def test_sanitize_attributes_redacts_content_keys_keeps_measures():
    out = sanitize_attributes(
        {
            "prompt": "what is attention?",
            "llm_response": "the full answer text",
            "retrieved_chunks": ["a", "b"],
            "raw_candidate_count": 12,
            "context_input_chars": 4096,
        }
    )
    assert out["prompt"] == REDACTED
    assert out["llm_response"] == REDACTED
    assert out["retrieved_chunks"] == REDACTED
    assert out["raw_candidate_count"] == 12
    assert out["context_input_chars"] == 4096


# --- Credentials / API headers ---


@pytest.mark.parametrize("key", ["api_key", "apikey", "password", "authorization", "secret", "bearer_token"])
def test_credential_keys_are_forbidden(key):
    assert is_forbidden_key(key) is True


def test_bearer_token_value_is_redacted():
    assert sanitize_attribute_value("Bearer abcdef123456ghijkl") == REDACTED


def test_openai_style_key_value_is_redacted():
    assert sanitize_attribute_value("sk-abcdEFGH1234567890ijklmnop") == REDACTED


def test_long_high_entropy_token_value_is_redacted():
    assert sanitize_attribute_value("A1B2C3D4E5F6G7H8I9J0K1L2M3N4O5P6Q7R8") == REDACTED


# --- Raw document IDs ---


def test_raw_document_id_key_is_forbidden():
    assert is_forbidden_key("document_id") is True
    assert is_forbidden_key("doc_id") is True


def test_sha256_document_id_value_is_redacted():
    sha = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
    assert sanitize_attribute_value(sha) == REDACTED


# --- Filesystem paths ---


@pytest.mark.parametrize(
    "value",
    [
        "/home/user/.studybuddy/pdfs/abc.pdf",
        "/usr/local/lib/python3.11/site-packages",
        r"C:\Users\SystemSu\.studybuddy\graphs\doc.json",
    ],
)
def test_filesystem_paths_are_redacted(value):
    assert sanitize_attribute_value(value) == REDACTED


# --- URL query strings ---


def test_url_query_string_is_stripped_but_path_kept():
    out = sanitize_attribute_value("http://127.0.0.1:4318/v1/traces?token=secret&x=1")
    assert out == "http://127.0.0.1:4318/v1/traces"


def test_url_fragment_is_stripped():
    out = sanitize_attribute_value("https://signoz.local/services#panel=abc")
    assert out == "https://signoz.local/services"


# --- base64 ---


def test_base64_blob_is_redacted():
    blob = "aGVsbG8gd29ybGQgdGhpcyBpcyBhIGxvbmcgYmFzZTY0IGJsb2IgcGF5bG9hZA=="
    assert sanitize_attribute_value(blob) == REDACTED


def test_data_uri_is_redacted():
    assert sanitize_attribute_value("data:image/png;base64,iVBORw0KGgoAAAANS") == REDACTED


# --- Oversized text ---


def test_oversized_string_is_truncated():
    # Realistic oversized metadata: a long label with whitespace (not an
    # opaque blob, which is redacted rather than truncated).
    value = "topic " * 1000
    out = sanitize_attribute_value(value)
    assert out != REDACTED
    assert len(out) < len(value)
    assert out.endswith("[truncated]")


def test_oversized_opaque_blob_is_redacted_not_truncated():
    # A very long whitespace-free base64-alphabet run is encoded content, so
    # redaction (never leaking a 256-char prefix) is the safe behaviour.
    assert sanitize_attribute_value("x" * 5000) == REDACTED


def test_short_plain_string_passes_through():
    assert sanitize_attribute_value("vector_search") == "vector_search"


# --- Scalars preserved ---


@pytest.mark.parametrize("value", [12, 3.14, True, False, 0])
def test_scalars_pass_through_unchanged(value):
    assert sanitize_attribute_value(value) == value


def test_none_is_dropped_from_attributes():
    out = sanitize_attributes({"a": None, "b": 1})
    assert "a" not in out
    assert out["b"] == 1


def test_nested_dict_value_is_dropped_to_avoid_content_leak():
    out = sanitize_attributes({"nested": {"prompt": "secret"}, "n": 2})
    assert "nested" not in out
    assert out["n"] == 2


# --- Exception text ---


def test_short_exception_text_kept():
    out = sanitize_exception_text("collection lookup failed")
    assert out == "collection lookup failed"


def test_sanitize_attributes_does_not_mutate_input():
    original = {"prompt": "x", "raw_candidate_count": 3}
    sanitize_attributes(original)
    assert original == {"prompt": "x", "raw_candidate_count": 3}