agent-lab/tests/test_benchmark_reporting.py

import importlib
import importlib.util
import json
from datetime import datetime, timezone
from pathlib import Path

import pytest
from pydantic import ValidationError

from agent_lab.application.benchmark import (
    BenchmarkCaseId,
    BenchmarkConfig,
    BenchmarkMode,
    BenchmarkRunResult,
)


def _reporting_module():
    spec = importlib.util.find_spec(
        "agent_lab.infrastructure.benchmark_reporting"
    )
    assert spec is not None, "benchmark_reporting module must be implemented"
    return importlib.import_module(
        "agent_lab.infrastructure.benchmark_reporting"
    )


def _config(**overrides: object) -> BenchmarkConfig:
    payload = {
        "schema_version": 1,
        "base_url": "https://provider.example/v1",
        "model": "benchmark-model",
        "runs_per_case": 1,
        "cases": [BenchmarkCaseId.ORDINARY_CHAT],
        "modes": [BenchmarkMode.DUAL_AGENT],
    }
    payload.update(overrides)
    return BenchmarkConfig.model_validate(payload)


def _run(
    *,
    case_id: BenchmarkCaseId = BenchmarkCaseId.ORDINARY_CHAT,
    mode: BenchmarkMode = BenchmarkMode.DUAL_AGENT,
    iteration: int = 1,
    status: str = "passed",
    initial_provider_ttft_ms: int | None = 10,
    visible_ttft_ms: int | None = 20,
    turn_wall_time_ms: int | None = 30,
    prompt_tokens: int | None = 2,
    completion_tokens: int | None = 3,
    total_tokens: int | None = 5,
    cached_tokens: int | None = 1,
    model_call_count: int | None = 1,
    fallback_count: int | None = 0,
    tool_count: int | None = 0,
    semantic_failures: list[str] | None = None,
    error: str | None = None,
) -> BenchmarkRunResult:
    return BenchmarkRunResult(
        case_id=case_id,
        mode=mode,
        iteration=iteration,
        status=status,
        initial_provider_ttft_ms=initial_provider_ttft_ms,
        visible_ttft_ms=visible_ttft_ms,
        turn_wall_time_ms=turn_wall_time_ms,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total_tokens,
        cached_tokens=cached_tokens,
        model_call_count=model_call_count,
        fallback_count=fallback_count,
        tool_count=tool_count,
        event_names=[],
        batch_event_names=[],
        tool_event_names=[],
        event_sources=[],
        tool_statuses=[],
        tool_latencies_ms=[],
        semantic_failures=semantic_failures or [],
        error=error,
        model_calls=[],
    )


@pytest.mark.parametrize(
    ("values", "quantile", "expected"),
    [
        ([], 0.50, None),
        ([7], 0.50, 7.0),
        ([7], 0.95, 7.0),
        ([1, 3, 5, 7], 0.50, 4.0),
        ([1, 3, 5, 7], 0.95, 6.7),
    ],
)
def test_percentile_uses_linear_interpolation(
    values: list[int], quantile: float, expected: float | None
):
    reporting = _reporting_module()

    assert reporting.percentile(values, quantile) == pytest.approx(expected)


def test_report_models_are_strict_and_include_required_summary_fields():
    reporting = _reporting_module()
    generated_at = datetime(2026, 7, 14, 1, 2, 3, tzinfo=timezone.utc)

    report = reporting.build_benchmark_report(
        _config(),
        [_run()],
        mock=True,
        generated_at=generated_at,
    )

    assert isinstance(report, reporting.BenchmarkReport)
    assert report.schema_version == 1
    assert report.generated_at == generated_at
    assert report.target.model_dump() == {
        "base_url": "https://provider.example/v1",
        "model": "benchmark-model",
        "mock": True,
    }
    assert report.config.model_dump(mode="json") == {
        "runs_per_case": 1,
        "cases": ["ordinary_chat"],
        "modes": ["dual_agent"],
    }
    assert report.overall.model_dump() == {
        "attempted": 1,
        "passed": 1,
        "failed": 0,
        "warnings": report.overall.warnings,
    }
    with pytest.raises(ValidationError):
        reporting.BenchmarkReport.model_validate(
            report.model_dump() | {"schema_version": "1"}
        )
    with pytest.raises(ValidationError):
        reporting.BenchmarkAggregate.model_validate(
            report.aggregates[0].model_dump() | {"unexpected": True}
        )


def test_aggregate_excludes_failed_and_null_samples_from_metrics():
    reporting = _reporting_module()
    runs = [
        _run(
            iteration=1,
            initial_provider_ttft_ms=10,
            visible_ttft_ms=20,
            turn_wall_time_ms=30,
            prompt_tokens=2,
            completion_tokens=3,
            total_tokens=5,
            cached_tokens=1,
            model_call_count=1,
            fallback_count=0,
            tool_count=2,
        ),
        _run(
            iteration=2,
            initial_provider_ttft_ms=None,
            visible_ttft_ms=40,
            turn_wall_time_ms=None,
            prompt_tokens=None,
            completion_tokens=7,
            total_tokens=7,
            cached_tokens=None,
            model_call_count=3,
            fallback_count=1,
            tool_count=None,
        ),
        _run(
            iteration=3,
            status="failed",
            initial_provider_ttft_ms=999,
            visible_ttft_ms=999,
            turn_wall_time_ms=999,
            prompt_tokens=999,
            completion_tokens=999,
            total_tokens=999,
            cached_tokens=999,
            model_call_count=999,
            fallback_count=999,
            tool_count=999,
            error="failed",
        ),
    ]

    aggregate = reporting.build_benchmark_report(
        _config(runs_per_case=3), runs, mock=True
    ).aggregates[0]

    assert (aggregate.attempted, aggregate.passed, aggregate.failed) == (3, 2, 1)
    assert aggregate.initial_provider_ttft_ms.model_dump() == {
        "count": 1,
        "p50": 10.0,
        "p95": 10.0,
    }
    assert aggregate.visible_ttft_ms.model_dump() == {
        "count": 2,
        "p50": 30.0,
        "p95": 39.0,
    }
    assert aggregate.turn_wall_time_ms.count == 1
    assert aggregate.prompt_tokens.model_dump() == {
        "count": 1,
        "total": 2,
        "mean": 2.0,
    }
    assert aggregate.completion_tokens.model_dump() == {
        "count": 2,
        "total": 10,
        "mean": 5.0,
    }
    assert aggregate.total_tokens.total == 12
    assert aggregate.cached_tokens.total == 1
    assert aggregate.model_call_count.model_dump() == {
        "count": 2,
        "total": 4,
        "mean": 2.0,
    }
    assert aggregate.fallback_count.total == 1
    assert aggregate.tool_count.model_dump() == {
        "count": 1,
        "total": 2,
        "mean": 2.0,
    }


def test_aggregates_are_stably_sorted_and_warn_for_low_confidence_p95():
    reporting = _reporting_module()
    runs = [
        _run(
            case_id=BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS,
            mode=BenchmarkMode.DUAL_AGENT,
        ),
        _run(
            case_id=BenchmarkCaseId.ORDINARY_CHAT,
            mode=BenchmarkMode.CHAT_AGENT_TOOLS,
        ),
        _run(
            case_id=BenchmarkCaseId.ORDINARY_CHAT,
            mode=BenchmarkMode.DUAL_AGENT,
        ),
    ]

    report = reporting.build_benchmark_report(
        _config(
            cases=[
                BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS,
                BenchmarkCaseId.ORDINARY_CHAT,
            ],
            modes=[
                BenchmarkMode.CHAT_AGENT_TOOLS,
                BenchmarkMode.DUAL_AGENT,
            ],
        ),
        runs,
        mock=True,
    )

    assert [(item.case_id.value, item.mode.value) for item in report.aggregates] == [
        ("ordinary_chat", "chat_agent_tools"),
        ("ordinary_chat", "dual_agent"),
        ("web_search_two_answers", "dual_agent"),
    ]
    assert all(
        any("p95 low confidence" in warning for warning in item.warnings)
        for item in report.aggregates
    )
    assert report.overall.warnings == [
        warning
        for aggregate in report.aggregates
        for warning in aggregate.warnings
    ]


def test_errors_are_redacted_before_truncation_in_json_and_markdown():
    reporting = _reporting_module()
    real_key = "sk-live-secret-value"
    leaked_error = (
        f"Authorization: Bearer {real_key}; "
        f"API-Key={real_key}; "
        f"https://example.test/?api_key={real_key}&token={real_key}; "
        f'{{"access_token":"{real_key}","key":"{real_key}"}}; '
        + "x" * 3000
    )

    report = reporting.build_benchmark_report(
        _config(),
        [
            _run(
                status="failed",
                semantic_failures=[f"token={real_key}"],
                error=leaked_error,
            )
        ],
        mock=False,
        secrets=[real_key],
    )
    payload = report.model_dump(mode="json")
    json_text = reporting.report_to_json(report)
    markdown = reporting.render_markdown(report)

    assert len(report.runs[0].error) == 2048
    assert real_key not in json.dumps(payload)
    assert real_key not in json_text
    assert real_key not in markdown
    assert "[REDACTED]" in json_text
    assert "[REDACTED]" in markdown


def test_configured_secret_is_redacted_from_all_report_string_surfaces():
    reporting = _reporting_module()
    real_key = "sk-everywhere-secret"
    run = _run().model_copy(
        update={
            "event_names": [real_key],
            "event_sources": [f"Bearer {real_key}"],
            "tool_statuses": [f"token={real_key}"],
        }
    )

    report = reporting.build_benchmark_report(
        _config(
            base_url=f"https://provider.example/{real_key}",
            model=f"model-{real_key}",
        ),
        [run],
        mock=False,
        secrets=[real_key],
    )
    rendered = reporting.report_to_json(report) + reporting.render_markdown(report)

    assert real_key not in rendered
    assert "[REDACTED]" in rendered


def test_markdown_is_rendered_from_the_same_json_facts():
    reporting = _reporting_module()
    report = reporting.build_benchmark_report(
        _config(runs_per_case=2),
        [
            _run(iteration=1, visible_ttft_ms=10, total_tokens=5),
            _run(iteration=2, visible_ttft_ms=30, total_tokens=9),
        ],
        mock=True,
        generated_at=datetime(2026, 7, 14, tzinfo=timezone.utc),
    )
    payload = json.loads(reporting.report_to_json(report))
    markdown = reporting.render_markdown(report)
    aggregate = payload["aggregates"][0]

    assert str(aggregate["visible_ttft_ms"]["p50"]) in markdown
    assert str(aggregate["visible_ttft_ms"]["p95"]) in markdown
    assert str(aggregate["total_tokens"]["total"]) in markdown
    assert str(aggregate["total_tokens"]["mean"]) in markdown
    assert payload == report.model_dump(mode="json")


def test_atomic_write_creates_matching_unique_json_and_markdown_paths(tmp_path: Path):
    reporting = _reporting_module()
    report = reporting.build_benchmark_report(
        _config(), [_run()], mock=True
    )
    timestamp = "20260714T010203Z"
    (tmp_path / f"{timestamp}.json").write_text("existing")
    (tmp_path / f"{timestamp}.md").write_text("existing")

    json_path, markdown_path = reporting.write_benchmark_report(
        report,
        tmp_path,
        timestamp=timestamp,
    )

    assert json_path.name == f"{timestamp}-1.json"
    assert markdown_path.name == f"{timestamp}-1.md"
    assert json.loads(json_path.read_text()) == report.model_dump(mode="json")
    assert markdown_path.read_text() == reporting.render_markdown(report)
    assert list(tmp_path.glob("*.tmp")) == []


def test_atomic_write_removes_temps_and_published_half_on_replace_failure(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
    reporting = _reporting_module()
    report = reporting.build_benchmark_report(
        _config(), [_run()], mock=True
    )
    real_replace = reporting.os.replace
    calls = 0

    def fail_second_replace(source: Path, destination: Path) -> None:
        nonlocal calls
        calls += 1
        if calls == 2:
            raise OSError("simulated markdown publish failure")
        real_replace(source, destination)

    monkeypatch.setattr(reporting.os, "replace", fail_second_replace)

    with pytest.raises(
        reporting.BenchmarkReportWriteError,
        match="failed to write benchmark reports",
    ):
        reporting.write_benchmark_report(
            report,
            tmp_path,
            timestamp="20260714T010203Z",
        )

    assert list(tmp_path.iterdir()) == []