test(proxy): cover tier keys and self-serve budget policies for key budget updates

tests/integration/authorization/test_key_budget_owner_updates.py

@@ -1,9 +1,11 @@
 from collections.abc import Iterator
 from hashlib import sha256
+from pathlib import Path
 from types import MappingProxyType
 from typing import Final
 
 import pytest
+import yaml
 from integration._support.client import (
     JSON_OBJECT,
     Gateway,
@@ -33,6 +35,30 @@ def resetting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
         yield replica
 
 
+def _policy_config(directory: Path, policy: str) -> Path:
+    shared: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    general_settings: Final = {
+        **JSON_OBJECT.validate_python(shared["general_settings"]),
+        "self_serve_budget_policy": policy,
+    }
+    path: Final = directory / f"{policy}.yaml"
+    path.write_text(yaml.safe_dump({**shared, "general_settings": general_settings}))
+    return path
+
+
+@pytest.fixture(scope="module", params=("lower_only", "ceiling"))
+def self_serve(
+    request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory
+) -> Iterator[tuple[str, Gateway]]:
+    policy: Final = string_value(request.param)
+    directory: Final = tmp_path_factory.mktemp(policy)
+    with (
+        gateway_from_environment() as shared,
+        owned_proxy(shared, directory, {}, config=_policy_config(directory, policy)) as replica,
+    ):
+        yield policy, replica
+
+
 def _owner_credential(scenario: Scenario) -> str:
     return scenario.key(user_id=scenario.user(user_role="internal_user"))
 
@@ -86,7 +112,7 @@ def test_key_owner_cannot_shorten_budget_duration_to_restart_an_exhausted_budget
         assert _chat_status(resetting, model, target) == BUDGET_EXCEEDED
 
 
-def test_key_owner_cannot_swap_budget_id_to_a_higher_tier(gateway: Gateway) -> None:
+def test_key_owner_cannot_swap_or_reschedule_a_tier_key(gateway: Gateway) -> None:
     with gateway.scenario() as scenario:
         model: Final = scenario.model()
         low: Final = scenario.budget(max_budget=CAP)
@@ -95,11 +121,14 @@ def test_key_owner_cannot_swap_budget_id_to_a_higher_tier(gateway: Gateway) -> N
         target: Final = _owned_key(gateway, scenario, owner, models=[model], budget_id=low)
         _exhaust(gateway, target)
         assert _chat_status(gateway, model, target) == BUDGET_EXCEEDED
+        before: Final = _key_row(gateway, target)
 
-        denied: Final = gateway.request("POST", "/key/update", {"key": target, "budget_id": high}, key=owner)
+        swapped: Final = gateway.request("POST", "/key/update", {"key": target, "budget_id": high}, key=owner)
+        rescheduled: Final = gateway.request("POST", "/key/update", {"key": target, "budget_duration": "1s"}, key=owner)
 
-        assert denied.status_code == 403, denied.text
-        assert _key_row(gateway, target)["budget_id"] == low
+        assert swapped.status_code == 403, swapped.text
+        assert rescheduled.status_code == 403, rescheduled.text
+        assert _key_row(gateway, target) == before
         assert _chat_status(gateway, model, target) == BUDGET_EXCEEDED
 
 
@@ -130,3 +159,37 @@ def test_key_owner_keeps_edits_that_cannot_restart_a_budget(gateway: Gateway) ->
         assert resent.status_code == 200, resent.text
         assert scheduled.status_code == 200, scheduled.text
         assert _key_row(gateway, uncapped)["budget_duration"] == "1d"
+
+
+def test_self_serve_policy_keeps_tier_keys_admin_only(self_serve: tuple[str, Gateway]) -> None:
+    _, proxy = self_serve
+    with proxy.scenario() as scenario:
+        low: Final = scenario.budget(max_budget=CAP)
+        high: Final = scenario.budget(max_budget=100.0)
+        owner: Final = _owner_credential(scenario)
+        target: Final = _owned_key(proxy, scenario, owner, budget_id=low)
+        before: Final = _key_row(proxy, target)
+
+        swapped: Final = proxy.request("POST", "/key/update", {"key": target, "budget_id": high}, key=owner)
+        rescheduled: Final = proxy.request("POST", "/key/update", {"key": target, "budget_duration": "1d"}, key=owner)
+
+        assert swapped.status_code == 403, swapped.text
+        assert rescheduled.status_code == 403, rescheduled.text
+        assert _key_row(proxy, target) == before
+
+
+def test_self_serve_policy_gates_a_new_reset_period_like_a_raised_cap(self_serve: tuple[str, Gateway]) -> None:
+    policy, proxy = self_serve
+    with proxy.scenario() as scenario:
+        owner: Final = _owner_credential(scenario)
+        shortened: Final = _owned_key(proxy, scenario, owner, max_budget=CAP, budget_duration="30d")
+        cleared: Final = _owned_key(proxy, scenario, owner, max_budget=CAP, budget_duration="30d")
+
+        shorten: Final = proxy.request("POST", "/key/update", {"key": shortened, "budget_duration": "1s"}, key=owner)
+        clear: Final = proxy.request("POST", "/key/update", {"key": cleared, "budget_duration": None}, key=owner)
+
+        assert (shorten.status_code, _key_row(proxy, shortened)["budget_duration"]) == (
+            (403, "30d") if policy == "lower_only" else (200, "1s")
+        ), shorten.text
+        assert clear.status_code == 200, clear.text
+        assert _key_row(proxy, cleared)["budget_duration"] is None