devplacepy/quota.py at d895de1b47ce5e5c6f1fa0385ac5f93da22612fe

 # retoor <retoor@molodetz.nl>
 from datetime import datetime, timedelta, timezone
 import pytest
 from pydantic import ValidationError
 from devplacepy.database import get_table
 from devplacepy.services.openai_gateway import quota as q
 from devplacepy.services.openai_gateway.usage import GATEWAY_LEDGER
 from devplacepy.utils import generate_uid
 _counter = [0]
 def _owner():
     _counter[0] += 1
     return f"quotauser{_counter[0]}-{generate_uid()}"
 def _burn(owner_id, app_reference, cost, minutes_ago=0):
     stamp = datetime.now(timezone.utc) - timedelta(minutes=minutes_ago)
     get_table(GATEWAY_LEDGER).insert(
         {
             "uid": generate_uid(),
             "owner_kind": "user",
             "owner_id": owner_id,
             "app_reference": app_reference,
             "cost_usd": cost,
             "created_at": stamp.isoformat(),
         }
     )
 def test_reset_requires_no_dimension(local_db):
     assert q.QuotaResetIn().owner_kind is None
     assert q.QuotaResetIn().owner_id is None
     assert q.QuotaResetIn().app_reference is None
 def test_reset_rejects_an_unknown_owner_kind(local_db):
     with pytest.raises(ValidationError):
         q.QuotaResetIn(owner_kind="wizard")
 def test_reset_rejects_a_malformed_app_reference(local_db):
     with pytest.raises(ValidationError):
         q.QuotaResetIn(app_reference="not a valid app!")
 def test_reset_normalizes_blanks_to_wildcards(local_db):
     payload = q.QuotaResetIn(owner_kind="", owner_id="  ", app_reference="")
     assert payload.owner_kind is None
     assert payload.owner_id is None
     assert payload.app_reference is None
 def test_spend_counts_before_any_reset(local_db):
     owner = _owner()
     _burn(owner, "appa", 1.5)
     assert q.spent_24h("user", owner, "appa") == 1.5
 def test_a_scoped_reset_clears_that_scope(local_db):
     owner = _owner()
     _burn(owner, "appa", 1.5)
     q.reset(q.QuotaResetIn(owner_kind="user", owner_id=owner, app_reference="appa"))
     assert q.spent_24h("user", owner, "appa") == 0.0
 def test_a_scoped_reset_leaves_another_caller_alone(local_db):
     first, second = _owner(), _owner()
     _burn(first, "appa", 1.5)
     _burn(second, "appa", 2.0)
     q.reset(q.QuotaResetIn(owner_kind="user", owner_id=first, app_reference="appa"))
     assert q.spent_24h("user", first, "appa") == 0.0
     assert q.spent_24h("user", second, "appa") == 2.0
 def test_a_scoped_reset_leaves_another_app_alone(local_db):
     owner = _owner()
     _burn(owner, "appa", 1.5)
     _burn(owner, "appb", 2.0)
     q.reset(q.QuotaResetIn(owner_kind="user", owner_id=owner, app_reference="appa"))
     assert q.spent_24h("user", owner, "appa") == 0.0
     assert q.spent_24h("user", owner, "appb") == 2.0
 def test_a_global_reset_clears_every_scope(local_db):
     first, second = _owner(), _owner()
     _burn(first, "appa", 1.5)
     _burn(second, "appb", 2.0)
     q.reset(q.QuotaResetIn())
     assert q.spent_24h("user", first, "appa") == 0.0
     assert q.spent_24h("user", second, "appb") == 0.0
 def test_spend_after_a_reset_counts_again(local_db):
     owner = _owner()
     _burn(owner, "appa", 1.5)
     q.reset(q.QuotaResetIn(owner_kind="user", owner_id=owner, app_reference="appa"))
     _burn(owner, "appa", 0.75)
     assert q.spent_24h("user", owner, "appa") == 0.75
 def test_a_reset_keeps_the_usage_history(local_db):
     owner = _owner()
     _burn(owner, "appa", 1.5)
     before = get_table(GATEWAY_LEDGER).count(owner_id=owner)
     q.reset(q.QuotaResetIn(owner_kind="user", owner_id=owner, app_reference="appa"))
     assert get_table(GATEWAY_LEDGER).count(owner_id=owner) == before
 def test_a_narrower_reset_does_not_clear_a_broader_scope(local_db):
     owner = _owner()
     _burn(owner, "appa", 1.5)
     q.reset(q.QuotaResetIn(owner_kind="user", owner_id=owner, app_reference="appa"))
     assert q.spent_24h("user", owner, None) == 1.5
 def test_a_broader_reset_clears_a_narrower_scope(local_db):
     owner = _owner()
     _burn(owner, "appa", 1.5)
     q.reset(q.QuotaResetIn(owner_kind="user", owner_id=owner))
     assert q.spent_24h("user", owner, "appa") == 0.0
 def test_resetting_the_same_scope_twice_reuses_one_row(local_db):
     owner = _owner()
     scope = q.QuotaResetIn(owner_kind="user", owner_id=owner, app_reference="appa")
     first = q.reset(scope)
     second = q.reset(scope)
     assert first["uid"] == second["uid"]
     assert second["reset_at"] >= first["reset_at"]
 def test_a_rule_scope_is_unblocked_by_a_reset(local_db):
     owner = _owner()
     q.quota_rule_store.set(
         q.QuotaRuleIn(
             owner_kind="user", owner_id=owner, app_reference="appa", limit_usd=1.0
         ),
         created_by="test",
     )
     _burn(owner, "appa", 1.5)
     limit, scope, rule = q.resolve("user", owner, "appa", {})
     assert q.spent_24h(*scope) >= limit
     q.reset(
         q.QuotaResetIn(
             owner_kind=scope[0], owner_id=scope[1], app_reference=scope[2]
         )
     )
     assert q.spent_24h(*scope) < limit
     q.quota_rule_store.remove(rule.uid)
 def test_scope_label_reads_like_the_rule_label(local_db):
     scope = {"owner_kind": "user", "owner_id": "u1", "app_reference": "appa"}
     assert q.scope_label(scope) == "role=user, user=u1, app=appa"
 def test_scope_label_falls_back_for_a_wildcard_scope(local_db):
     empty = {"owner_kind": None, "owner_id": None, "app_reference": None}
     assert q.scope_label(empty, fallback="every caller") == "every caller"