Aether/tests/services/test_stats_aggregator_optimization.py

from __future__ import annotations

from datetime import date, datetime, time, timedelta, timezone
from types import SimpleNamespace
from typing import Any, cast

import pytest
from sqlalchemy.exc import IntegrityError

from src.models.database import StatsDaily, StatsUserDaily
from src.services.system.stats_aggregator import (
    AggregatedStats,
    StatsAggregatorService,
    query_stats_hybrid,
)
from src.services.system.time_range import TimeRangeParams


class _FakeQuery:
    def __init__(self, *, all_result: list[Any] | None = None) -> None:
        self._all_result = all_result if all_result is not None else []

    def filter(self, *_args: object, **_kwargs: object) -> _FakeQuery:
        return self

    def group_by(self, *_args: object, **_kwargs: object) -> _FakeQuery:
        return self

    def all(self) -> list[Any]:
        return self._all_result


class _HybridQuerySession:
    def __init__(self, stats_daily_rows: list[SimpleNamespace]) -> None:
        self._stats_daily_rows = stats_daily_rows
        self.stats_daily_query_count = 0

    def query(self, entity: object) -> _FakeQuery:
        if entity is StatsDaily:
            self.stats_daily_query_count += 1
            return _FakeQuery(all_result=self._stats_daily_rows)
        raise AssertionError(f"Unexpected query entity: {entity}")


class _BatchUserStatsSession:
    def __init__(
        self, existing_rows: list[StatsUserDaily], aggregated_rows: list[SimpleNamespace]
    ) -> None:
        self._responses: list[list[Any]] = [list(existing_rows), list(aggregated_rows)]
        self.added: list[StatsUserDaily] = []
        self.commit_count = 0

    def query(self, *_entities: object) -> _FakeQuery:
        if not self._responses:
            raise AssertionError("Unexpected extra query")
        return _FakeQuery(all_result=self._responses.pop(0))

    def add(self, row: StatsUserDaily) -> None:
        self.added.append(row)

    def commit(self) -> None:
        self.commit_count += 1


class _RetryCommitSession:
    def __init__(self) -> None:
        self.commit_count = 0
        self.rollback_count = 0

    def commit(self) -> None:
        self.commit_count += 1
        if self.commit_count == 1:
            raise IntegrityError("insert", {}, Exception("duplicate key"))

    def rollback(self) -> None:
        self.rollback_count += 1


def test_query_stats_hybrid_batches_statsdaily_lookup_and_merges_realtime_ranges(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    today = datetime.now(timezone.utc).date()
    historical_cached_day = today - timedelta(days=4)
    historical_missing_day = today - timedelta(days=3)
    realtime_day = today

    cached_row = SimpleNamespace(
        date=datetime.combine(historical_cached_day, time.min, tzinfo=timezone.utc),
        total_requests=10,
        success_requests=9,
        error_requests=1,
        input_tokens=100,
        output_tokens=50,
        cache_creation_tokens=5,
        cache_read_tokens=3,
        cache_creation_cost=1.2,
        cache_read_cost=0.8,
        total_cost=3.5,
        actual_total_cost=3.0,
        avg_response_time_ms=200.0,
    )
    db = _HybridQuerySession(stats_daily_rows=[cached_row])

    calls: list[tuple[datetime, datetime]] = []

    def _fake_aggregate_usage_range(
        _db: object,
        start_utc: datetime,
        end_utc: datetime,
        filters: object | None = None,  # noqa: ARG001
    ) -> AggregatedStats:
        calls.append((start_utc, end_utc))
        return AggregatedStats(total_requests=1, success_requests=1)

    class _FakeParams:
        def get_complete_utc_dates(self) -> tuple[list[date], None, None]:
            return [historical_cached_day, historical_missing_day, realtime_day], None, None

    monkeypatch.setattr(
        "src.services.system.stats_aggregator.aggregate_usage_range",
        _fake_aggregate_usage_range,
    )

    result = query_stats_hybrid(cast(Any, db), cast(Any, _FakeParams()))

    assert db.stats_daily_query_count == 1
    assert calls == [
        (
            datetime.combine(historical_missing_day, time.min, tzinfo=timezone.utc),
            datetime.combine(
                historical_missing_day + timedelta(days=1), time.min, tzinfo=timezone.utc
            ),
        ),
        (
            datetime.combine(realtime_day, time.min, tzinfo=timezone.utc),
            datetime.combine(realtime_day + timedelta(days=1), time.min, tzinfo=timezone.utc),
        ),
    ]
    assert result.total_requests == 12
    assert result.success_requests == 11


def test_aggregate_user_daily_stats_batch_updates_all_users_in_two_queries() -> None:
    target_day = datetime(2026, 3, 1, tzinfo=timezone.utc)
    aggregated_rows = [
        SimpleNamespace(
            user_id="user-1",
            username="alice",
            total_requests=4,
            error_requests=1,
            input_tokens=20,
            output_tokens=8,
            cache_creation_tokens=2,
            cache_read_tokens=1,
            total_cost=1.5,
        )
    ]
    db = _BatchUserStatsSession(existing_rows=[], aggregated_rows=aggregated_rows)

    result = StatsAggregatorService.aggregate_user_daily_stats_batch(
        cast(Any, db),
        target_day,
        ["user-1", "user-2"],
        commit=True,
    )

    assert len(result) == 2
    assert db.commit_count == 1
    assert len(db.added) == 2

    user_one = next(row for row in result if row.user_id == "user-1")
    user_two = next(row for row in result if row.user_id == "user-2")

    assert user_one.username == "alice"
    assert user_one.total_requests == 4
    assert user_one.success_requests == 3
    assert user_one.total_cost == 1.5

    assert user_two.total_requests == 0
    assert user_two.success_requests == 0
    assert user_two.error_requests == 0
    assert user_two.total_cost == 0.0


def test_aggregate_daily_stats_bundle_retries_after_integrity_error(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    target_day = datetime(2026, 3, 1, tzinfo=timezone.utc)
    db = _RetryCommitSession()
    stats_calls: list[SimpleNamespace] = []
    stage_calls: list[str] = []

    def fake_aggregate_daily_stats(*_args: object, **_kwargs: object) -> SimpleNamespace:
        stats = SimpleNamespace(is_complete=False, aggregated_at=None)
        stats_calls.append(stats)
        stage_calls.append("daily")
        return stats

    def fake_model_stats(*_args: object, **_kwargs: object) -> list[object]:
        stage_calls.append("model")
        return []

    def fake_provider_stats(*_args: object, **_kwargs: object) -> list[object]:
        stage_calls.append("provider")
        return []

    def fake_api_key_stats(*_args: object, **_kwargs: object) -> list[object]:
        stage_calls.append("api_key")
        return []

    def fake_error_stats(*_args: object, **_kwargs: object) -> list[object]:
        stage_calls.append("error")
        return []

    def fake_user_daily_stats(*_args: object, **_kwargs: object) -> list[object]:
        stage_calls.append("user")
        return []

    monkeypatch.setattr(StatsAggregatorService, "aggregate_daily_stats", fake_aggregate_daily_stats)
    monkeypatch.setattr(StatsAggregatorService, "aggregate_daily_model_stats", fake_model_stats)
    monkeypatch.setattr(
        StatsAggregatorService, "aggregate_daily_provider_stats", fake_provider_stats
    )
    monkeypatch.setattr(StatsAggregatorService, "aggregate_daily_api_key_stats", fake_api_key_stats)
    monkeypatch.setattr(StatsAggregatorService, "aggregate_daily_error_stats", fake_error_stats)
    monkeypatch.setattr(
        StatsAggregatorService, "aggregate_user_daily_stats_batch", fake_user_daily_stats
    )

    result = StatsAggregatorService.aggregate_daily_stats_bundle(
        cast(Any, db),
        target_day,
        user_ids=["user-1"],
    )

    assert db.commit_count == 2
    assert db.rollback_count == 1
    assert len(stats_calls) == 2
    assert stage_calls == [
        "daily",
        "model",
        "provider",
        "api_key",
        "error",
        "user",
        "daily",
        "model",
        "provider",
        "api_key",
        "error",
        "user",
    ]
    assert result is stats_calls[-1]
    assert result.is_complete is True
    assert result.aggregated_at is not None


def test_compute_percentiles_by_local_day_returns_sqlite_fallback_without_queries() -> None:
    db = SimpleNamespace(bind=SimpleNamespace(dialect=SimpleNamespace(name="sqlite")))
    time_range = TimeRangeParams(
        start_date=date(2026, 3, 1),
        end_date=date(2026, 3, 3),
        timezone="Asia/Singapore",
    )

    result = StatsAggregatorService.compute_percentiles_by_local_day(cast(Any, db), time_range)

    assert [row["date"] for row in result] == ["2026-03-01", "2026-03-02", "2026-03-03"]
    assert all(row["p50_response_time_ms"] is None for row in result)
    assert all(row["p50_first_byte_time_ms"] is None for row in result)