update

AzA march 2026/tests/test_ai_budget_phase1.py

@@ -0,0 +1,170 @@
+# -*- coding: utf-8 -*-
+"""Lokale Unit-Tests KI-Budget Phase 1 (kein Netzwerk)."""
+from __future__ import annotations
+
+import os
+import sqlite3
+import tempfile
+import unittest
+from pathlib import Path
+
+from aza_ai_budget import (
+    LicenseBudgetRow,
+    check_allows_openai_call,
+    compute_budget_snapshot,
+    ensure_ai_budget_schema,
+    estimate_openai_cost_usd,
+    insert_usage_event,
+    resolve_license_for_device,
+    sum_usage_usd_for_period,
+)
+
+
+def _mk_db(path: Path) -> None:
+    con = sqlite3.connect(str(path))
+    con.execute(
+        """
+        CREATE TABLE device_bindings (
+          id INTEGER PRIMARY KEY AUTOINCREMENT,
+          customer_email TEXT NOT NULL,
+          user_key TEXT NOT NULL,
+          device_hash TEXT NOT NULL,
+          first_seen_at INTEGER NOT NULL,
+          last_seen_at INTEGER NOT NULL,
+          is_active INTEGER DEFAULT 1,
+          UNIQUE(customer_email, user_key, device_hash)
+        )
+        """
+    )
+    con.execute(
+        """
+        CREATE TABLE licenses (
+          subscription_id TEXT PRIMARY KEY,
+          customer_id TEXT,
+          status TEXT,
+          lookup_key TEXT,
+          allowed_users INTEGER,
+          devices_per_user INTEGER,
+          customer_email TEXT,
+          client_reference_id TEXT,
+          current_period_end INTEGER,
+          current_period_start INTEGER,
+          updated_at INTEGER NOT NULL,
+          license_key TEXT,
+          practice_id TEXT
+        )
+        """
+    )
+    now = 1_700_000_000
+    con.execute(
+        """
+        INSERT INTO licenses(subscription_id, customer_id, status, lookup_key, allowed_users, devices_per_user,
+          customer_email, client_reference_id, current_period_end, current_period_start, updated_at, license_key, practice_id)
+        VALUES ('sub_t1', 'cus_x', 'active', 'aza_basic_monthly', 1, 2,
+          'cli@example.test', NULL, ?, ?, ?, 'KEY', 'prac_test')
+        """,
+        (now + 86400 * 30, now, now),
+    )
+    import hashlib
+
+    dh = hashlib.sha256(b"device-unit-test").hexdigest()
+    con.execute(
+        """
+        INSERT INTO device_bindings(customer_email, user_key, device_hash, first_seen_at, last_seen_at)
+        VALUES ('cli@example.test', 'prac_test', ?, ?, ?)
+        """,
+        (dh, now, now),
+    )
+    con.commit()
+    con.close()
+
+
+class TestAiBudgetPhase1(unittest.TestCase):
+    def setUp(self):
+        self.tmp = tempfile.NamedTemporaryFile(suffix=".sqlite", delete=False)
+        self.tmp.close()
+        self.db_path = Path(self.tmp.name)
+        _mk_db(self.db_path)
+
+    def tearDown(self):
+        try:
+            os.unlink(self.db_path)
+        except OSError:
+            pass
+
+    def test_schema_idempotent(self):
+        with sqlite3.connect(str(self.db_path)) as con:
+            ensure_ai_budget_schema(con)
+            ensure_ai_budget_schema(con)
+            n = con.execute(
+                "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='ai_usage_events'"
+            ).fetchone()[0]
+            self.assertEqual(n, 1)
+
+    def test_resolve_and_percent(self):
+        with sqlite3.connect(str(self.db_path)) as con:
+            ensure_ai_budget_schema(con)
+            lic = resolve_license_for_device(con, "device-unit-test")
+            self.assertIsNotNone(lic)
+            snap = compute_budget_snapshot(con, lic)
+            self.assertEqual(snap["available_percent"], 100)
+            ok, _ = check_allows_openai_call(con, lic)
+            self.assertTrue(ok)
+
+    def test_usage_reduces_percent(self):
+        with sqlite3.connect(str(self.db_path)) as con:
+            ensure_ai_budget_schema(con)
+            lic = resolve_license_for_device(con, "device-unit-test")
+            ps, pe = lic.period_start or 0, lic.period_end or 0
+            insert_usage_event(
+                con,
+                lic=lic,
+                device_id="device-unit-test",
+                period_start=ps,
+                period_end=pe,
+                operation_type="chat",
+                model="gpt-4o-mini",
+                input_tokens=1_000_000,
+                output_tokens=0,
+                total_tokens=1_000_000,
+                audio_seconds=0.0,
+                estimated_cost_usd=estimate_openai_cost_usd(
+                    model="gpt-4o-mini", input_tokens=1_000_000, output_tokens=0
+                ),
+                request_id="t1",
+                status="success",
+            )
+            used = sum_usage_usd_for_period(con, lic.subscription_id, ps, pe)
+            self.assertGreater(used, 0)
+            snap2 = compute_budget_snapshot(con, lic)
+            self.assertLess(snap2["available_percent"], 100)
+
+    def test_block_at_budget(self):
+        with sqlite3.connect(str(self.db_path)) as con:
+            ensure_ai_budget_schema(con)
+            lic = resolve_license_for_device(con, "device-unit-test")
+            ps, pe = int(lic.period_start or 0), int(lic.period_end or 0)
+            big = 50.0
+            insert_usage_event(
+                con,
+                lic=lic,
+                device_id="x",
+                period_start=ps,
+                period_end=pe,
+                operation_type="chat",
+                model="gpt-4o-mini",
+                input_tokens=0,
+                output_tokens=0,
+                total_tokens=0,
+                audio_seconds=0.0,
+                estimated_cost_usd=big,
+                request_id="t2",
+                status="success",
+            )
+            ok, info = check_allows_openai_call(con, lic)
+            self.assertFalse(ok)
+            self.assertEqual(info.get("error_code"), "AI_BUDGET_EXCEEDED")
+
+
+if __name__ == "__main__":
+    unittest.main()