60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
"""`rag/budget.py::RequestBudget` — F-08."""
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from rag.budget import RequestBudget, RequestBudgetExhausted
|
|
from rag.ports import AnswerGenerationUnavailable
|
|
|
|
|
|
def test_budget_exhausted_is_a_subclass_of_answer_generation_unavailable():
|
|
"""Deliberate: every existing `except AnswerGenerationUnavailable:`
|
|
fail-open/fail-closed handler in the codebase must catch this with zero
|
|
changes, since it predates F-08 and already encodes the right behaviour
|
|
for "the provider is unavailable to us right now"."""
|
|
assert issubclass(RequestBudgetExhausted, AnswerGenerationUnavailable)
|
|
|
|
|
|
def test_fresh_budget_has_budget():
|
|
budget = RequestBudget.start(max_wall_clock_ms=20_000, max_calls=5)
|
|
assert budget.has_budget() is True
|
|
|
|
|
|
def test_require_spends_one_call():
|
|
budget = RequestBudget.start(max_wall_clock_ms=20_000, max_calls=2)
|
|
budget.require()
|
|
assert budget.calls_remaining == 1
|
|
budget.require()
|
|
assert budget.calls_remaining == 0
|
|
assert budget.has_budget() is False
|
|
|
|
|
|
def test_require_raises_once_calls_are_exhausted():
|
|
budget = RequestBudget.start(max_wall_clock_ms=20_000, max_calls=1)
|
|
budget.require()
|
|
with pytest.raises(RequestBudgetExhausted):
|
|
budget.require()
|
|
|
|
|
|
def test_require_raises_once_the_deadline_has_passed():
|
|
budget = RequestBudget.start(max_wall_clock_ms=0, max_calls=100)
|
|
time.sleep(0.01)
|
|
assert budget.has_budget() is False
|
|
with pytest.raises(RequestBudgetExhausted):
|
|
budget.require()
|
|
|
|
|
|
def test_a_failed_require_does_not_spend_a_call():
|
|
"""`require()` raises before decrementing when there's no budget left —
|
|
`calls_remaining` must not go negative, which would otherwise make a
|
|
budget that's already exhausted look like it has "negative debt" instead
|
|
of cleanly `0`."""
|
|
budget = RequestBudget.start(max_wall_clock_ms=20_000, max_calls=1)
|
|
budget.require()
|
|
for _ in range(3):
|
|
with pytest.raises(RequestBudgetExhausted):
|
|
budget.require()
|
|
assert budget.calls_remaining == 0
|