aza/AzA march 2026/_live_test_set_state.py

#!/usr/bin/env python3
"""Set marked test usage to fully exhaust monthly + extra credit for auto-topup gate test."""
import json
import sqlite3
import sys
import time
import uuid
from pathlib import Path

DB = Path("/root/aza-app/data/stripe_webhook.sqlite")
PRACTICE = "prac_883ddc21fb6a"
SUB = "wc_sub_1502"
TEST_REQUEST_ID = "test_popup_budget_20260521_exhaust"
sys.path.insert(0, "/root/aza-app")
from aza_ai_budget import LicenseBudgetRow, check_allows_openai_call, compute_budget_snapshot
from aza_ai_credit import compute_extra_credit_remaining


def load_license(con):
    row = con.execute(
        "SELECT * FROM licenses WHERE subscription_id = ? AND practice_id = ?",
        (SUB, PRACTICE),
    ).fetchone()
    cols = [d[1] for d in con.execute("PRAGMA table_info(licenses)").fetchall()]
    return dict(zip(cols, row))


def main():
    con = sqlite3.connect(str(DB))
    lic_d = load_license(con)
    ps = int(lic_d["current_period_start"])
    pe = int(lic_d["current_period_end"])
    lic = LicenseBudgetRow(
        subscription_id=lic_d["subscription_id"],
        customer_email=lic_d.get("customer_email") or "",
        customer_id=lic_d.get("customer_id"),
        practice_id=lic_d.get("practice_id"),
        lookup_key=lic_d.get("lookup_key"),
        status=lic_d.get("status") or "",
        period_start=ps,
        period_end=pe,
    )
    budget = 20.0
    extra = compute_extra_credit_remaining(
        con, practice_id=PRACTICE, subscription_id=SUB, monthly_budget_usd=budget
    )
    used = con.execute(
        """
        SELECT COALESCE(SUM(estimated_cost_usd), 0)
        FROM ai_usage_events
        WHERE subscription_id = ? AND status = 'success'
          AND period_start = ? AND period_end = ?
          AND request_id != ?
        """,
        (SUB, ps, pe, TEST_REQUEST_ID),
    ).fetchone()[0]
    # Exhaust monthly (20) + extra credit buffer
    target = round(float(used) + max(0.0, budget - float(used)) + max(0.0, extra) + 0.5, 4)
    delta = round(max(0.0, target - float(used)), 4)
    if delta <= 0:
        print("SET_SKIP delta<=0")
    else:
        con.execute(
            """
            INSERT INTO ai_usage_events(
              id, created_at, license_email, customer_id, subscription_id, practice_id,
              device_id_suffix, period_start, period_end, operation_type, model,
              input_tokens, output_tokens, total_tokens, audio_seconds, estimated_cost_usd,
              request_id, status, error_code, meta_json
            ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
            """,
            (
                str(uuid.uuid4()),
                int(time.time()),
                lic_d.get("customer_email"),
                lic_d.get("customer_id"),
                SUB,
                PRACTICE,
                "test0000",
                ps,
                pe,
                "test_popup_budget",
                "test-marker",
                0,
                0,
                0,
                0.0,
                delta,
                TEST_REQUEST_ID,
                "success",
                None,
                json.dumps({"test": "auto_topup_live_trigger", "temporary": True}, ensure_ascii=False),
            ),
        )
        con.commit()
        print(f"SET_OK delta_usd={delta} target={target} extra_before={extra}")
    snap = compute_budget_snapshot(con, lic)
    ok, _ = check_allows_openai_call(con, lic)
    print("SNAPSHOT=" + json.dumps({
        "monthly_budget_usd": snap.get("budget_usd"),
        "used_usd": snap.get("used_usd"),
        "monthly_remaining_usd": snap.get("remaining_usd"),
        "extra_credit_remaining_usd": snap.get("extra_credit_remaining_usd"),
        "total_available_usd": snap.get("total_available_usd"),
        "available_percent": snap.get("available_percent"),
        "allows_call": ok,
    }, ensure_ascii=False))


if __name__ == "__main__":
    main()