diff --git a/tests/integration/authorization/test_key_budget_owner_updates.py b/tests/integration/authorization/test_key_budget_owner_updates.py index f6e85aac5cc..8415c044b6f 100644 --- a/tests/integration/authorization/test_key_budget_owner_updates.py +++ b/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