from __future__ import annotations import builtins import sys from dataclasses import dataclass from datetime import datetime, timedelta from types import SimpleNamespace import pytest import headroom.reporting as reporting from headroom.reporting import generator @dataclass class FakeMetrics: request_id: str model: str mode: str timestamp: datetime tokens_input_before: int tokens_input_after: int cache_alignment_score: float waste_signals: dict[str, int] class FakeStorage: def __init__(self, stats: dict, items: list[FakeMetrics]) -> None: self._stats = stats self._items = items self.closed = False def get_summary_stats(self, start_time, end_time): return dict(self._stats) def iter_all(self): return iter(self._items) def close(self) -> None: self.closed = True def test_reporting_public_export() -> None: assert reporting.generate_report is generator.generate_report assert reporting.__all__ == ["generate_report"] def test_get_jinja2_template_success_with_stub(monkeypatch) -> None: class FakeTemplate: def __init__(self, template_str: str) -> None: self.template_str = template_str def render(self, **kwargs) -> str: return f"{self.template_str}:{kwargs['name']}" monkeypatch.setitem(sys.modules, "jinja2", SimpleNamespace(Template=FakeTemplate)) template = generator._get_jinja2_template("hello") assert template.render(name="world") == "hello:world" def test_get_jinja2_template_raises_helpful_error(monkeypatch) -> None: real_import = builtins.__import__ def fake_import(name, globals=None, locals=None, fromlist=(), level=0): if name == "jinja2": raise ImportError("missing") return real_import(name, globals, locals, fromlist, level) monkeypatch.setattr(builtins, "__import__", fake_import) with pytest.raises(ImportError, match="jinja2 is required for report generation"): generator._get_jinja2_template("ignored") def test_build_waste_histogram_empty_and_filtered_data() -> None: now = datetime(2026, 4, 23, 12, 0, 0) metrics = [ FakeMetrics( request_id="before", model="gpt-4o", mode="audit", timestamp=now - timedelta(days=2), tokens_input_before=100, tokens_input_after=90, cache_alignment_score=10, waste_signals={"json_bloat": 5}, ), FakeMetrics( request_id="inside", model="gpt-4o", mode="optimize", timestamp=now, tokens_input_before=200, tokens_input_after=100, cache_alignment_score=70, waste_signals={"json_bloat": 30, "html_noise": 10, "dynamic_date": 5, "reread": 20}, ), FakeMetrics( request_id="flat", model="gpt-4o", mode="audit", timestamp=now, tokens_input_before=50, tokens_input_after=50, cache_alignment_score=50, waste_signals={"whitespace": 4}, ), FakeMetrics( request_id="after", model="gpt-4o", mode="audit", timestamp=now + timedelta(days=2), tokens_input_before=100, tokens_input_after=20, cache_alignment_score=20, waste_signals={"base64": 50}, ), ] histogram = generator._build_waste_histogram( FakeStorage({}, metrics), start_time=now - timedelta(hours=1), end_time=now + timedelta(hours=1), ) assert histogram[0] == {"label": "History Bloat", "tokens": 55, "percentage": 100.0} assert histogram[1] == pytest.approx( {"label": "Tool JSON Bloat", "tokens": 30, "percentage": 54.54545454545454} ) # "reread" surfaces in the histogram but is excluded from known_waste, # so History Bloat above stays 100 - 45 = 55. assert any( item["label"] == "Re-served Tool Results" and item["tokens"] == 20 for item in histogram ) assert any(item["label"] == "HTML Noise" and item["tokens"] == 10 for item in histogram) assert any(item["label"] == "Dynamic Dates" and item["tokens"] == 5 for item in histogram) assert any(item["label"] == "Base64 Blobs" and item["tokens"] == 0 for item in histogram) empty = generator._build_waste_histogram(FakeStorage({}, []), None, None) assert all(item["tokens"] == 0 and item["percentage"] == 0 for item in empty) def test_get_top_waste_requests_sorts_filters_and_limits() -> None: now = datetime(2026, 4, 23, 12, 0, 0) metrics = [ FakeMetrics("one", "gpt-4o", "audit", now, 400, 100, 80, {}), FakeMetrics("two", "gpt-4o-mini", "optimize", now, 350, 330, 70, {}), FakeMetrics("three", "claude", "audit", now - timedelta(days=3), 1000, 10, 50, {}), FakeMetrics("four", "claude", "audit", now + timedelta(days=3), 1000, 200, 40, {}), ] top_requests = generator._get_top_waste_requests( FakeStorage({}, metrics), start_time=now - timedelta(hours=1), end_time=now + timedelta(hours=1), limit=1, ) assert top_requests == [ { "request_id": "one", "model": "gpt-4o", "mode": "audit", "tokens_before": 400, "tokens_saved": 300, "cache_alignment": 80, } ] def test_generate_recommendations_for_heavy_waste_and_for_getting_started() -> None: stats = { "avg_cache_alignment": 40, "audit_count": 7, "optimize_count": 3, "total_tokens_saved": 120000, "estimated_savings": "$1.23", } histogram = [ {"label": "Tool JSON Bloat", "tokens": 15000, "percentage": 100}, {"label": "History Bloat", "tokens": 60000, "percentage": 50}, ] recommendations = generator._generate_recommendations(stats, histogram, top_requests=[{}]) titles = [item["title"] for item in recommendations] assert titles == [ "Improve Cache Alignment", "Enable Tool Output Compression", "Review Rolling Window Settings", "Switch to Optimize Mode", "Continue Monitoring", ] assert "15,000" in recommendations[1]["description"] assert "60,000" in recommendations[2]["description"] starter = generator._generate_recommendations( { "avg_cache_alignment": 90, "audit_count": 1, "optimize_count": 1, "total_tokens_saved": 0, "estimated_savings": "$0.00", }, [{"label": "Tool JSON Bloat", "tokens": 1, "percentage": 100}], top_requests=[], ) assert starter == [ { "title": "Get Started", "description": "No optimizations applied yet. Try setting headroom_mode='optimize' " "on your next request to start seeing token savings.", } ] @pytest.mark.parametrize( ("start_time", "end_time", "expected_period"), [ ( datetime(2026, 4, 20, 8, 0, 0), datetime(2026, 4, 23, 18, 0, 0), "2026-04-20 to 2026-04-23", ), (datetime(2026, 4, 20, 8, 0, 0), None, "Since 2026-04-20"), (None, datetime(2026, 4, 23, 18, 0, 0), "Until 2026-04-23"), (None, None, "All time"), ], ) def test_generate_report_writes_output_and_closes_storage( monkeypatch, tmp_path, start_time, end_time, expected_period ) -> None: storage = FakeStorage( { "total_requests": 3, "total_tokens_saved": 50, "avg_tokens_saved": 16.6, "total_tokens_before": 100, "total_tokens_after": 0, "avg_cache_alignment": 82, "audit_count": 1, "optimize_count": 2, }, [], ) render_calls: list[dict] = [] class FakeTemplate: def render(self, **kwargs) -> str: render_calls.append(kwargs) return "report" monkeypatch.setattr(generator, "create_storage", lambda store_url: storage) monkeypatch.setattr( generator, "_build_waste_histogram", lambda *args: [{"label": "x", "tokens": 1}] ) monkeypatch.setattr( generator, "_get_top_waste_requests", lambda *args, **kwargs: [{"request_id": "abc"}] ) monkeypatch.setattr( generator, "_generate_recommendations", lambda *args: [{"title": "Keep going"}] ) monkeypatch.setattr(generator, "_get_jinja2_template", lambda template_str: FakeTemplate()) monkeypatch.setattr( generator, "estimate_cost", lambda tokens, output_tokens, model: {100: 2.0, 0: None}[tokens], ) monkeypatch.setattr(generator, "format_cost", lambda cost: f"${cost:.2f}") output_path = tmp_path / "report.html" result = generator.generate_report( "sqlite:///demo.db", output_path=str(output_path), start_time=start_time, end_time=end_time, ) assert result == str(output_path) assert output_path.read_text() == "report" assert render_calls[0]["period"] == expected_period assert render_calls[0]["stats"]["tpm_multiplier"] == 100.0 assert render_calls[0]["stats"]["estimated_savings"] == "$2.00" assert storage.closed is True def test_generate_report_closes_storage_when_render_fails(monkeypatch, tmp_path) -> None: storage = FakeStorage( { "total_requests": 0, "total_tokens_saved": 0, "avg_tokens_saved": 0, "total_tokens_before": 0, "total_tokens_after": 0, "avg_cache_alignment": 0, "audit_count": 0, "optimize_count": 0, }, [], ) class FakeTemplate: def render(self, **kwargs) -> str: raise RuntimeError("boom") monkeypatch.setattr(generator, "create_storage", lambda store_url: storage) monkeypatch.setattr(generator, "_build_waste_histogram", lambda *args: []) monkeypatch.setattr(generator, "_get_top_waste_requests", lambda *args, **kwargs: []) monkeypatch.setattr(generator, "_generate_recommendations", lambda *args: []) monkeypatch.setattr(generator, "_get_jinja2_template", lambda template_str: FakeTemplate()) monkeypatch.setattr(generator, "estimate_cost", lambda *args: 0.0) monkeypatch.setattr(generator, "format_cost", lambda cost: "$0.00") with pytest.raises(RuntimeError, match="boom"): generator.generate_report("sqlite:///demo.db", output_path=str(tmp_path / "report.html")) assert storage.closed is True