tests/test_usage_tracker.py

import json
import os
from datetime import datetime

import psycopg2
import pytest
import tempfile
import time
from app.core.usage_tracker import UsageTracker, Tier, UsageRecord


@pytest.fixture
def tracker():
    with tempfile.NamedTemporaryFile(suffix=".db") as tmp:
        yield UsageTracker(db_path=tmp.name)


def test_get_or_create_api_key(tracker):
    # Updated: pass tenant_id as keyword argument (new signature)
    assert tracker.get_or_create_api_key("test_key", tenant_id="test") is True
    assert tracker.get_tier("test_key") == Tier.FREE
    # Second call should return True without error
    assert tracker.get_or_create_api_key("test_key", tenant_id="test") is True


def test_update_api_key_tier(tracker):
    tracker.get_or_create_api_key("test_key", tenant_id="test")
    assert tracker.update_api_key_tier("test_key", Tier.PRO) is True
    assert tracker.get_tier("test_key") == Tier.PRO
    # Non-existent key
    assert tracker.update_api_key_tier("nonexistent", Tier.PRO) is False


def test_get_remaining_quota_free(tracker):
    tracker.get_or_create_api_key("free_key", tenant_id="test")
    # Initially 1000 remaining
    remaining = tracker.get_remaining_quota("free_key", Tier.FREE)
    assert remaining == 1000
    # Simulate usage using the atomic method
    record = UsageRecord(
        api_key="free_key",
        tier=Tier.FREE,
        timestamp=time.time(),
        endpoint="/test"
    )
    tracker.increment_usage_sync(record)
    remaining = tracker.get_remaining_quota("free_key", Tier.FREE)
    assert remaining == 999


def test_get_remaining_quota_enterprise(tracker):
    tracker.get_or_create_api_key("ent_key", tenant_id="test")
    remaining = tracker.get_remaining_quota("ent_key", Tier.ENTERPRISE)
    assert remaining is None


def test_increment_usage_sync(tracker):
    tracker.get_or_create_api_key("test_key", tenant_id="test")
    record = UsageRecord(
        api_key="test_key",
        tier=Tier.FREE,
        timestamp=time.time(),
        endpoint="/test",
    )
    result = tracker.increment_usage_sync(record)
    assert result is True
    # Check quota decreased
    remaining = tracker.get_remaining_quota("test_key", Tier.FREE)
    assert remaining == 999


def test_get_audit_logs(tracker):
    tracker.get_or_create_api_key("test_key", tenant_id="test")
    record = UsageRecord(
        api_key="test_key",
        tier=Tier.FREE,
        timestamp=time.time(),
        endpoint="/test",
        request_body={"foo": "bar"},
        response={"status": "ok"},
    )
    tracker.increment_usage_sync(record)
    logs = tracker.get_audit_logs("test_key", limit=10)
    assert len(logs) == 1
    assert logs[0]["endpoint"] == "/test"


def _pg_monthly_count(api_key: str, month: str) -> int:
    """Direct Postgres read, bypassing UsageTracker entirely -- this is
    exactly the query arf-gateway's Go code runs against the same table."""
    conn = psycopg2.connect(os.environ["DATABASE_URL"])
    try:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT COALESCE(count, 0) FROM monthly_counts WHERE api_key = %s AND year_month = %s",
                (api_key, month),
            )
            row = cur.fetchone()
            return row[0] if row else 0
    finally:
        conn.close()


def test_increment_usage_sync_mirrors_to_postgres_monthly_counts(tracker):
    """Regresses arf-gateway-001: arf-gateway's quota check reads
    monthly_counts from Postgres, not from this service's local SQLite
    file. Every successfully-counted call must be visible there."""
    tracker.get_or_create_api_key("pg-mirror-key", tenant_id="test")
    month = tracker._get_month_key()

    # No truncate fixture exists for this Postgres-resident table (CI's
    # Postgres service is ephemeral per run, but a local repeat run against
    # a persistent database could have leftover rows) -- assert the delta,
    # not an absolute count.
    baseline = _pg_monthly_count("pg-mirror-key", month)

    record = UsageRecord(
        api_key="pg-mirror-key", tier=Tier.FREE, timestamp=time.time(), endpoint="/test",
    )
    tracker.increment_usage_sync(record)
    assert _pg_monthly_count("pg-mirror-key", month) == baseline + 1

    tracker.increment_usage_sync(record)
    assert _pg_monthly_count("pg-mirror-key", month) == baseline + 2


def test_increment_usage_sync_succeeds_even_if_postgres_mirror_fails(tracker, monkeypatch):
    """The local quota decision (SQLite/Redis) must not fail just because
    the best-effort mirror write to Postgres did -- see
    _record_pg_monthly_count's docstring. A degraded mirror write should
    not turn into a degraded evaluate/healing endpoint for the caller."""
    def _boom(*args, **kwargs):
        raise psycopg2.OperationalError("simulated Postgres outage")

    # get_or_create_api_key also goes through _get_pg_conn (it persists to
    # the real api_keys table), so it must run before the patch below -- the
    # outage this test simulates is specific to the monthly_counts mirror
    # write, not to Postgres as a whole.
    tracker.get_or_create_api_key("mirror-fail-key", tenant_id="test")

    # Patch what _record_pg_monthly_count calls internally, not the method
    # itself -- replacing the whole method would bypass its own try/except
    # and prove nothing about that error-handling actually working.
    monkeypatch.setattr(tracker, "_get_pg_conn", _boom)

    record = UsageRecord(
        api_key="mirror-fail-key", tier=Tier.FREE, timestamp=time.time(), endpoint="/test",
    )

    # Must not raise, and the local quota decision must still succeed.
    result = tracker.increment_usage_sync(record)
    assert result is True
    assert tracker.get_remaining_quota("mirror-fail-key", Tier.FREE) == 999


def test_insert_audit_log_writes_response_row(tracker):
    """routes_governance.py schedules current_tracker._insert_audit_log as
    a background task (background_tasks.add_task) to record the response
    body once it's known, at app/api/routes_governance.py:407 and :738 --
    always with tier=None, since quota was already consumed by an earlier
    consume_quota_and_log call for the same request. The real UsageTracker
    had no such method (only tests/conftest.py's MockTracker did), so every
    real call raised AttributeError inside the background task (arf-api-002)."""
    record = UsageRecord(
        api_key="audit-log-key",
        tier=None,
        timestamp=time.time(),
        endpoint="/api/v1/intents/evaluate/response",
        request_body=None,
        response={"recommended_action": "approve"},
        processing_ms=12.5,
    )

    tracker._insert_audit_log(record)

    logs = tracker.get_audit_logs("audit-log-key", limit=10)
    assert len(logs) == 1
    assert logs[0]["endpoint"] == "/api/v1/intents/evaluate/response"
    assert logs[0]["tier"] == "unknown"
    assert json.loads(logs[0]["response"]) == {"recommended_action": "approve"}


def test_consume_quota_and_log_handles_non_json_native_request_body(tracker):
    """request_body/response are whatever a Pydantic model's plain
    .model_dump() returns, e.g. ReliabilityEvent.timestamp in
    app/api/routes_governance.py's HealingDecisionRequest -- a raw
    datetime, not the ISO string model_dump(mode="json") would produce.
    json.dumps has no default encoder for datetime, so any real request
    carrying one raised TypeError here, outside any try/except in the
    /healing/evaluate handler, on every call (surfaced while verifying the
    tier=None fix for that same endpoint: fixing tier alone still crashed,
    one layer deeper, on this). default=str makes the insert tolerant of
    datetime and any other type json.dumps doesn't natively handle."""
    tracker.get_or_create_api_key("datetime-body-key", tenant_id="test")
    record = UsageRecord(
        api_key="datetime-body-key",
        tier=Tier.FREE,
        timestamp=time.time(),
        endpoint="/api/v1/healing/evaluate",
        request_body={"event": {"component": "svc", "timestamp": datetime.now()}},
    )

    result = tracker.increment_usage_sync(record)
    assert result is True

    logs = tracker.get_audit_logs("datetime-body-key", limit=10)
    assert len(logs) == 1
    assert "timestamp" in json.loads(logs[0]["request_body"])["event"]