From 2d57a62467ff8903b01d7128eba6a93c3748217c Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Mon, 2 Nov 2026 11:26:51 +0100 Subject: [PATCH 1/3] LCORE-4610: add unit tests for the quota scheduler and token history The quota scheduler and the token usage history had no direct tests: nothing failed if a limiter type mapped to the wrong subject, if the scheduler started without storage, or if the PostgreSQL statements took different parameters from the SQLite ones. Cover the subject mapping, the choice of statements, connection checks, the startup conditions and reconnection attempts of the scheduler, and the parameters of the PostgreSQL statements, which no test can run against a real server here. For the token usage history cover the running totals per user, provider and model, reconnecting after a closed connection, running without storage, and a write failing loudly when its table is gone. --- tests/unit/quota/README.md | 4 + tests/unit/quota/test_token_usage_history.py | 98 ++++++ tests/unit/runners/README.md | 4 + tests/unit/runners/test_quota_scheduler.py | 318 +++++++++++++++++++ 4 files changed, 424 insertions(+) create mode 100644 tests/unit/quota/test_token_usage_history.py create mode 100644 tests/unit/runners/test_quota_scheduler.py diff --git a/tests/unit/quota/README.md b/tests/unit/quota/README.md index 8bc76d7fd..26f2535b7 100644 --- a/tests/unit/quota/README.md +++ b/tests/unit/quota/README.md @@ -24,6 +24,10 @@ Unit tests for QuotaExceedError class. Unit tests for quota limiter factory class. +## [test_token_usage_history.py](test_token_usage_history.py) + +Unit tests for the TokenUsageHistory class. + ## [test_user_quota_limiter.py](test_user_quota_limiter.py) Unit tests for UserQuotaLimiter class. diff --git a/tests/unit/quota/test_token_usage_history.py b/tests/unit/quota/test_token_usage_history.py new file mode 100644 index 000000000..6d31f97d4 --- /dev/null +++ b/tests/unit/quota/test_token_usage_history.py @@ -0,0 +1,98 @@ +"""Unit tests for the TokenUsageHistory class.""" + +import sqlite3 + +import pytest + +from models.config import QuotaHandlersConfiguration, SQLiteDatabaseConfiguration +from quota.token_usage_history import TokenUsageHistory + +USER = "user-1" +OTHER_USER = "user-2" +PROVIDER = "test-provider" +MODEL = "test-model" + + +@pytest.fixture(name="history") +def history_fixture() -> TokenUsageHistory: + """Create a token usage history on an in-memory SQLite database.""" + configuration = QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + sqlite=SQLiteDatabaseConfiguration(db_path=":memory:"), + enable_token_history=True, + ) + return TokenUsageHistory(configuration) + + +def usage_rows(history: TokenUsageHistory) -> dict[tuple[str, str, str], tuple]: + """Read every token_usage row of the history, keyed by its primary key.""" + assert history.connection is not None + cursor = history.connection.cursor() + cursor.execute( + "SELECT user_id, provider, model, input_tokens, output_tokens FROM token_usage" + ) + rows = cursor.fetchall() + cursor.close() + return {(row[0], row[1], row[2]): (row[3], row[4]) for row in rows} + + +def test_connected_with_storage(history: TokenUsageHistory) -> None: + """A history created on a storage is connected and has its table.""" + assert history.connected() + assert not usage_rows(history) + + +def test_usage_accumulates_per_user_provider_and_model( + history: TokenUsageHistory, +) -> None: + """Every (user, provider, model) triple keeps its own running totals.""" + history.consume_tokens(USER, PROVIDER, MODEL, 100, 50) + history.consume_tokens(USER, PROVIDER, MODEL, 20, 5) + history.consume_tokens(USER, PROVIDER, "other-model", 1, 1) + history.consume_tokens(OTHER_USER, PROVIDER, MODEL, 7, 3) + + assert usage_rows(history) == { + (USER, PROVIDER, MODEL): (120, 55), + (USER, PROVIDER, "other-model"): (1, 1), + (OTHER_USER, PROVIDER, MODEL): (7, 3), + } + + +def test_reconnects_after_the_connection_was_closed( + history: TokenUsageHistory, +) -> None: + """A dead connection is reopened before the next write.""" + assert history.connection is not None + history.connection.close() + assert not history.connected() + + history.consume_tokens(USER, PROVIDER, MODEL, 1, 2) + + assert history.connected() + # :memory: databases do not survive their connection, so the reopened + # history starts empty and holds only the write made after the reconnect + assert usage_rows(history) == {(USER, PROVIDER, MODEL): (1, 2)} + + +def test_without_storage_nothing_is_connected_and_writes_are_ignored() -> None: + """Without a storage the history has no connection and consuming is a no-op.""" + history = TokenUsageHistory( + QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + enable_token_history=True + ) + ) + + assert history.connection is None + assert not history.connected() + history.consume_tokens(USER, PROVIDER, MODEL, 1, 1) + assert history.connection is None + + +def test_write_fails_loudly_when_the_table_is_gone( + history: TokenUsageHistory, +) -> None: + """A storage error is not swallowed; the caller maps it to an HTTP 500.""" + assert history.connection is not None + history.connection.execute("DROP TABLE token_usage") + + with pytest.raises(sqlite3.OperationalError): + history.consume_tokens(USER, PROVIDER, MODEL, 1, 1) diff --git a/tests/unit/runners/README.md b/tests/unit/runners/README.md index 92086f296..2933665a9 100644 --- a/tests/unit/runners/README.md +++ b/tests/unit/runners/README.md @@ -4,6 +4,10 @@ Unit tests for runners. +## [test_quota_scheduler.py](test_quota_scheduler.py) + +Unit tests for the quota scheduler runner. + ## [test_uvicorn_runner.py](test_uvicorn_runner.py) Unit tests for the Uvicorn runner implementation. diff --git a/tests/unit/runners/test_quota_scheduler.py b/tests/unit/runners/test_quota_scheduler.py new file mode 100644 index 000000000..222b73326 --- /dev/null +++ b/tests/unit/runners/test_quota_scheduler.py @@ -0,0 +1,318 @@ +"""Unit tests for the quota scheduler runner. + +These tests cover the parts of the scheduler that need no database: the +statement selection per backend, the subject mapping, the early returns of +``quota_scheduler`` and the shape of the PostgreSQL statements. The renewal +itself, on a real SQLite database, is covered by +``tests/integration/test_quota_scheduler.py``. +""" + +import sqlite3 +from typing import Any + +import pytest +from pydantic import SecretStr +from pytest_mock import MockerFixture + +import constants +from models.config import ( + PostgreSQLDatabaseConfiguration, + QuotaHandlersConfiguration, + QuotaLimiterConfiguration, + QuotaSchedulerConfiguration, + SQLiteDatabaseConfiguration, +) +from quota.sql import ( + CREATE_QUOTA_TABLE_PG, + INCREASE_QUOTA_STATEMENT_PG, + INCREASE_QUOTA_STATEMENT_SQLITE, + INIT_QUOTA_PG, + INIT_QUOTA_SQLITE, + RESET_QUOTA_STATEMENT_PG, + RESET_QUOTA_STATEMENT_SQLITE, + SELECT_QUOTA_PG, + SELECT_QUOTA_SQLITE, + SET_AVAILABLE_QUOTA_PG, + SET_AVAILABLE_QUOTA_SQLITE, + UPDATE_AVAILABLE_QUOTA_PG, + UPDATE_AVAILABLE_QUOTA_SQLITE, +) +from runners import quota_scheduler +from runners.quota_scheduler import ( + connect, + connected, + get_increase_quota_statement, + get_reset_quota_statement, + get_subject_id, + increase_quota, + init_tables, + quota_revocation, + reset_quota, +) + +SCHEDULER_MODULE = "runners.quota_scheduler" + +POSTGRES = PostgreSQLDatabaseConfiguration( # pyright: ignore[reportCallIssue] + db="quota", + user="quota", + password=SecretStr("secret"), + host="localhost", + port=5432, +) +SQLITE = SQLiteDatabaseConfiguration(db_path=":memory:") + + +def user_limiter( + initial_quota: int = 100, quota_increase: int = 10, period: str = "1 day" +) -> QuotaLimiterConfiguration: + """Build the configuration of one user limiter.""" + return QuotaLimiterConfiguration( + type="user_limiter", + name="user quota", + initial_quota=initial_quota, + quota_increase=quota_increase, + period=period, + ) + + +# ========================================== +# Pure helpers +# ========================================== + + +@pytest.mark.parametrize( + ("limiter_type", "expected"), + [ + (constants.USER_QUOTA_LIMITER, "u"), + (constants.CLUSTER_QUOTA_LIMITER, "c"), + ("something_else", "?"), + ], +) +def test_get_subject_id(limiter_type: str, expected: str) -> None: + """A limiter type maps to the subject column of the quota table.""" + assert get_subject_id(limiter_type) == expected + + +@pytest.mark.parametrize( + ("config", "expected_increase", "expected_reset"), + [ + pytest.param( + QuotaHandlersConfiguration( + sqlite=SQLITE + ), # pyright: ignore[reportCallIssue] + INCREASE_QUOTA_STATEMENT_SQLITE, + RESET_QUOTA_STATEMENT_SQLITE, + id="sqlite", + ), + pytest.param( + QuotaHandlersConfiguration( + postgres=POSTGRES + ), # pyright: ignore[reportCallIssue] + INCREASE_QUOTA_STATEMENT_PG, + RESET_QUOTA_STATEMENT_PG, + id="postgres", + ), + pytest.param( + QuotaHandlersConfiguration( + sqlite=SQLITE, postgres=POSTGRES + ), # pyright: ignore[reportCallIssue] + INCREASE_QUOTA_STATEMENT_SQLITE, + RESET_QUOTA_STATEMENT_SQLITE, + id="both_prefers_sqlite", + ), + ], +) +def test_statement_selection( + config: QuotaHandlersConfiguration, expected_increase: str, expected_reset: str +) -> None: + """The scheduler picks the statements of the configured backend.""" + assert get_increase_quota_statement(config) == expected_increase + assert get_reset_quota_statement(config) == expected_reset + + +def test_connected_on_live_closed_and_missing_connection() -> None: + """``connected`` answers for a live, a closed and a missing connection.""" + connection = sqlite3.connect(":memory:") + assert connected(connection) is True + connection.close() + assert connected(connection) is False + assert connected(None) is False + + +# ========================================== +# quota_scheduler early returns +# ========================================== + + +@pytest.mark.parametrize( + "config", + [ + pytest.param( + QuotaHandlersConfiguration(), # pyright: ignore[reportCallIssue] + id="no_storage", + ), + pytest.param( + QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + limiters=[user_limiter()] + ), + id="limiters_without_storage", + ), + pytest.param( + QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + sqlite=SQLITE + ), + id="no_limiters", + ), + ], +) +def test_quota_scheduler_does_not_start_without_storage_or_limiters( + config: QuotaHandlersConfiguration, mocker: MockerFixture +) -> None: + """Without storage or limiters the scheduler returns without connecting.""" + connect_mock = mocker.patch(f"{SCHEDULER_MODULE}.connect") + + assert quota_scheduler.quota_scheduler(config) is False + + connect_mock.assert_not_called() + + +def test_quota_scheduler_gives_up_after_the_configured_reconnection_attempts( + mocker: MockerFixture, +) -> None: + """An unreachable database is retried the configured number of times, then skipped.""" + config = QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + sqlite=SQLiteDatabaseConfiguration(db_path="/nonexistent/dir/quota.db"), + limiters=[user_limiter()], + scheduler=QuotaSchedulerConfiguration( + period=1, database_reconnection_count=3, database_reconnection_delay=7 + ), + ) + connect_sqlite = mocker.patch( + f"{SCHEDULER_MODULE}.connect_sqlite", + side_effect=sqlite3.OperationalError("unable to open database file"), + ) + sleep = mocker.patch(f"{SCHEDULER_MODULE}.sleep") + + assert quota_scheduler.quota_scheduler(config) is False + + assert connect_sqlite.call_count == 3 + assert sleep.call_args_list == [mocker.call(7)] * 3 + + +# ========================================== +# PostgreSQL: statement shape, with a mocked connection +# ========================================== + +# No PostgreSQL service exists in CI or locally, so the PostgreSQL path is +# checked for shape only: the statements take the same parameters in the same +# order as the SQLite ones that the integration tests execute for real. + + +@pytest.mark.parametrize( + ("pg_statement", "sqlite_statement"), + [ + pytest.param(INCREASE_QUOTA_STATEMENT_PG, INCREASE_QUOTA_STATEMENT_SQLITE), + pytest.param(RESET_QUOTA_STATEMENT_PG, RESET_QUOTA_STATEMENT_SQLITE), + pytest.param(INIT_QUOTA_PG, INIT_QUOTA_SQLITE), + pytest.param(SELECT_QUOTA_PG, SELECT_QUOTA_SQLITE), + pytest.param(SET_AVAILABLE_QUOTA_PG, SET_AVAILABLE_QUOTA_SQLITE), + pytest.param(UPDATE_AVAILABLE_QUOTA_PG, UPDATE_AVAILABLE_QUOTA_SQLITE), + ], + ids=["increase", "reset", "init", "select", "set_available", "update_available"], +) +def test_pg_and_sqlite_statements_take_the_same_parameters( + pg_statement: str, sqlite_statement: str +) -> None: + """Both backends are driven by one parameter tuple, so the counts must match.""" + assert pg_statement.count("%s") == sqlite_statement.count("?") + + +def _pg_connection(mocker: MockerFixture) -> Any: + """Build a DB-API connection double that records what is executed on it.""" + connection = mocker.Mock() + connection.cursor.return_value.rowcount = 1 + return connection + + +def test_increase_and_reset_execute_pg_statements_with_value_subject_period( + mocker: MockerFixture, +) -> None: + """The PostgreSQL statements get (value, subject, period), then a commit.""" + connection = _pg_connection(mocker) + + increase_quota(connection, INCREASE_QUOTA_STATEMENT_PG, "u", 10, "1 day") + reset_quota(connection, RESET_QUOTA_STATEMENT_PG, "c", 100, "1 month") + + cursor = connection.cursor.return_value + assert cursor.execute.call_args_list == [ + mocker.call(INCREASE_QUOTA_STATEMENT_PG, (10, "u", "1 day")), + mocker.call(RESET_QUOTA_STATEMENT_PG, (100, "c", "1 month")), + ] + assert connection.commit.call_count == 2 + + +def test_quota_revocation_runs_both_pg_statements_for_the_limiter_subject( + mocker: MockerFixture, +) -> None: + """With PostgreSQL configured, revocation issues the PG increase and reset. + + Which of the two runs first is not pinned: the order decides what a row + with both values set ends up with (see the integration tests) and is + being reported, not asserted. + """ + config = QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + postgres=POSTGRES, + limiters=[user_limiter(initial_quota=100, quota_increase=10, period="1 day")], + ) + connection = _pg_connection(mocker) + + quota_revocation( + connection, + config.limiters[0], + get_increase_quota_statement(config), + get_reset_quota_statement(config), + ) + + cursor = connection.cursor.return_value + assert cursor.execute.call_count == 2 + cursor.execute.assert_any_call(INCREASE_QUOTA_STATEMENT_PG, (10, "u", "1 day")) + cursor.execute.assert_any_call(RESET_QUOTA_STATEMENT_PG, (100, "u", "1 day")) + + +def test_init_tables_executes_the_create_statement_and_commits( + mocker: MockerFixture, +) -> None: + """Table creation runs the given statement once and commits.""" + connection = _pg_connection(mocker) + + init_tables(connection, CREATE_QUOTA_TABLE_PG) + + connection.cursor.return_value.execute.assert_called_once_with( + CREATE_QUOTA_TABLE_PG + ) + connection.commit.assert_called_once() + + +def test_connect_prefers_postgres_and_falls_back_to_sqlite( + mocker: MockerFixture, +) -> None: + """``connect`` opens PostgreSQL when configured, SQLite otherwise, nothing without both.""" + connect_pg = mocker.patch(f"{SCHEDULER_MODULE}.connect_pg") + connect_sqlite = mocker.patch(f"{SCHEDULER_MODULE}.connect_sqlite") + + both = QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + postgres=POSTGRES, sqlite=SQLITE + ) + sqlite_only = QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + sqlite=SQLITE + ) + neither = QuotaHandlersConfiguration() # pyright: ignore[reportCallIssue] + + assert connect(both) is connect_pg.return_value + connect_pg.assert_called_once_with(POSTGRES) + connect_sqlite.assert_not_called() + + assert connect(sqlite_only) is connect_sqlite.return_value + connect_sqlite.assert_called_once_with(SQLITE) + + assert connect(neither) is None From 5370b95e413cd2e421035c5acdd2f6a76c9eca18 Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Mon, 2 Nov 2026 12:43:37 +0100 Subject: [PATCH 2/3] LCORE-4610: test quota renewal on a real SQLite database LCORE-4606 showed that the renewal statements were never executed by any test: on SQLite a quota was renewed at every scheduler run. These tests run the real statements and scheduler functions on a real database. They pin that a quota is renewed only once its period has elapsed and once per period, for both the reset and the increase; that the time a limiter stamps on a new row is read correctly by SQLite; that a zero value skips its statement; that a renewal touches only the rows of its own subject; and that the scheduler loop creates the table, survives a failing limiter and reconnects after its connection died. With initial_quota and quota_increase both set, the increase runs first and stamps revoked_at, so the reset of the same run no longer matches and the balance is not what the quota_limiter module documents. That is pinned as a strict xfail rather than fixed here. --- tests/integration/_quota_helpers.py | 193 ++++++++++ tests/integration/test_quota_scheduler.py | 447 ++++++++++++++++++++++ 2 files changed, 640 insertions(+) create mode 100644 tests/integration/_quota_helpers.py create mode 100644 tests/integration/test_quota_scheduler.py diff --git a/tests/integration/_quota_helpers.py b/tests/integration/_quota_helpers.py new file mode 100644 index 000000000..178934b60 --- /dev/null +++ b/tests/integration/_quota_helpers.py @@ -0,0 +1,193 @@ +"""Shared helpers for the quota management integration tests (LCORE-4610). + +The quota limiters and the quota scheduler share one SQLite table, +``quota_limits``, with one row per (subject id, subject type). The helpers +here build limiter configurations, read rows straight from the database and +move the ``revoked_at`` stamp of rows into the past, so that a renewal can be +observed without waiting for a period to elapse. +""" + +import sqlite3 +from collections.abc import Iterator +from contextlib import contextmanager +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any, Literal, Optional + +from models.config import ( + QuotaHandlersConfiguration, + QuotaLimiterConfiguration, + QuotaSchedulerConfiguration, + SQLiteDatabaseConfiguration, +) +from quota.quota_limiter import QuotaLimiter +from runners.quota_scheduler import ( + get_increase_quota_statement, + get_reset_quota_statement, + quota_revocation, +) + +USER_SUBJECT = "u" +CLUSTER_SUBJECT = "c" +CLUSTER = "" +"""The cluster limiter keeps one row, under the empty id.""" + +ROW_AGE_SECONDS = 3 +"""Rows are back-dated by this much: past a 2 second period, inside a 10 second one.""" + +TimestampStyle = Literal["limiter", "scheduler"] + + +def limiter_config( + limiter_type: str, + *, + initial_quota: int, + quota_increase: int, + period: str, +) -> QuotaLimiterConfiguration: + """Build the configuration of one limiter.""" + return QuotaLimiterConfiguration( + type=limiter_type, # type: ignore[arg-type] + name=f"{limiter_type} test", + initial_quota=initial_quota, + quota_increase=quota_increase, + period=period, + ) + + +def handlers( + db_path: Path, + *limiters: QuotaLimiterConfiguration, + scheduler_period: int = 1, + enable_token_history: bool = False, +) -> QuotaHandlersConfiguration: + """Build a quota configuration on the given SQLite file.""" + return QuotaHandlersConfiguration( # pyright: ignore[reportCallIssue] + sqlite=SQLiteDatabaseConfiguration(db_path=str(db_path)), + limiters=list(limiters), + scheduler=QuotaSchedulerConfiguration( + period=scheduler_period + ), # pyright: ignore[reportCallIssue] + enable_token_history=enable_token_history, + ) + + +@contextmanager +def _connected(db_path: Path) -> Iterator[sqlite3.Connection]: + """Open a connection for one helper call, commit on success, close always. + + The ``sqlite3`` connection as a context manager only wraps a transaction; + it never closes, and the endpoint tests call the helpers dozens of times + on a file the limiters and the scheduler keep open. + """ + connection = sqlite3.connect(db_path) + try: + with connection: + yield connection + finally: + connection.close() + + +def read_available(db_path: Path, row_id: str, subject: str) -> Optional[int]: + """Read the balance of one row straight from the database, None if no row.""" + with _connected(db_path) as connection: + row = connection.execute( + "SELECT available FROM quota_limits WHERE id=? AND subject=?", + (row_id, subject), + ).fetchone() + return None if row is None else int(row[0]) + + +def read_revoked_at(db_path: Path, row_id: str, subject: str) -> str: + """Read the ``revoked_at`` stamp of one row as the text SQLite stores.""" + with _connected(db_path) as connection: + (revoked_at,) = connection.execute( + "SELECT revoked_at FROM quota_limits WHERE id=? AND subject=?", + (row_id, subject), + ).fetchone() + return str(revoked_at) + + +def read_row_age_seconds(db_path: Path, row_id: str, subject: str) -> Optional[float]: + """Read how long ago SQLite thinks one row was last renewed. + + The age is computed by SQLite itself from the ``revoked_at`` text, the way + the renewal statements read it. It is None when the stamp is missing or + is in a format ``datetime()`` can not parse, which is the case a renewal + silently never matches. + """ + with _connected(db_path) as connection: + (age,) = connection.execute( + "SELECT (julianday('now') - julianday(revoked_at)) * 86400" + " FROM quota_limits WHERE id=? AND subject=?", + (row_id, subject), + ).fetchone() + return None if age is None else float(age) + + +def table_exists(db_path: Path, table: str) -> bool: + """Tell whether the database file holds a table of the given name.""" + with _connected(db_path) as connection: + row = connection.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,) + ).fetchone() + return row is not None + + +def read_token_usage(db_path: Path) -> dict[tuple[str, str, str], tuple[int, int]]: + """Read the token usage history, keyed by (user, provider, model).""" + with _connected(db_path) as connection: + rows = connection.execute( + "SELECT user_id, provider, model, input_tokens, output_tokens" + " FROM token_usage" + ).fetchall() + return {(row[0], row[1], row[2]): (int(row[3]), int(row[4])) for row in rows} + + +def back_date( + db_path: Path, + seconds: int = ROW_AGE_SECONDS, + style: TimestampStyle = "limiter", + subject: Optional[str] = None, +) -> None: + """Move the ``revoked_at`` stamp of the rows into the past. + + ``limiter`` writes the stamp as the limiter does, an ISO timestamp with a + ``T`` separator; ``scheduler`` writes it as the renewal statements do, with + SQLite's ``datetime()`` text. Without ``subject`` every row is moved. + """ + where = "" if subject is None else " WHERE subject=?" + params: tuple[Any, ...] = () if subject is None else (subject,) + with _connected(db_path) as connection: + if style == "limiter": + stamp = datetime.now(tz=UTC) - timedelta(seconds=seconds) + connection.execute( + f"UPDATE quota_limits SET revoked_at=?{where}", + (stamp.replace(tzinfo=None).isoformat(), *params), + ) + else: + connection.execute( + f"UPDATE quota_limits SET revoked_at=datetime('now', ?){where}", + (f"-{seconds} seconds", *params), + ) + + +def spend(limiter: QuotaLimiter, tokens: int, subject: str = CLUSTER) -> None: + """Use a limiter the way a request does: read the balance first, then consume. + + The read creates the row of a subject the database has not seen yet, with + ``initial_quota`` and the current time as its ``revoked_at`` stamp. + """ + limiter.available_quota(subject) + limiter.consume_tokens(tokens, 0, subject) + + +def renew(connection: Any, config: QuotaHandlersConfiguration) -> None: + """Run one scheduler sync: the renewal of every configured limiter.""" + for limiter in config.limiters: + quota_revocation( + connection, + limiter, + get_increase_quota_statement(config), + get_reset_quota_statement(config), + ) diff --git a/tests/integration/test_quota_scheduler.py b/tests/integration/test_quota_scheduler.py new file mode 100644 index 000000000..e5f06c8b3 --- /dev/null +++ b/tests/integration/test_quota_scheduler.py @@ -0,0 +1,447 @@ +"""Integration tests for the periodic quota renewal (LCORE-4610). + +The quota scheduler renews the balances the limiters keep in the quota +database: once a limiter's period has elapsed for a row, the row is reset to +``initial_quota`` and/or increased by ``quota_increase``, and its +``revoked_at`` stamp moves so that the next renewal waits for a whole period +again (LCORE-4606). + +Everything here runs on a real SQLite file: the rows are created and consumed +through the real limiters, and the renewal runs the scheduler's own +statements over the scheduler's own connection. Time is not waited for; the +``revoked_at`` stamp of a row is moved into the past instead, in both formats +found in a real database (the limiter writes an ISO ``T`` timestamp, the +scheduler rewrites it with SQLite's ``datetime('now')``). The stamp the +limiter really writes is checked once, as SQLite reads it. +""" + +import sqlite3 +from collections.abc import Callable, Generator +from pathlib import Path +from typing import Any + +import pytest +from pytest_mock import MockerFixture + +from models.config import QuotaHandlersConfiguration, QuotaLimiterConfiguration +from quota.cluster_quota_limiter import ClusterQuotaLimiter +from quota.quota_limiter import QuotaLimiter +from quota.sql import CREATE_QUOTA_TABLE_SQLITE +from quota.user_quota_limiter import UserQuotaLimiter +from runners import quota_scheduler +from runners.quota_scheduler import connect +from tests.integration._quota_helpers import ( + CLUSTER, + CLUSTER_SUBJECT, + USER_SUBJECT, + back_date, + handlers, + limiter_config, + read_available, + read_revoked_at, + read_row_age_seconds, + renew, + spend, + table_exists, +) + +ALICE, BOB = "alice", "bob" + +ELAPSED = "2 seconds" +NOT_ELAPSED = "10 seconds" +PAST_NOT_ELAPSED = 11 +"""Back-dating by this much puts a row past a ``NOT_ELAPSED`` period.""" + + +class StopLoop(Exception): + """Raised from the patched ``sleep`` to leave the scheduler's endless loop.""" + + +@pytest.fixture(name="quota_db") +def quota_db_fixture(tmp_path: Path) -> Path: + """Path of the SQLite quota database of the test.""" + return tmp_path / "quota.db" + + +@pytest.fixture(name="make_limiter") +def make_limiter_fixture() -> Generator[Callable[..., QuotaLimiter], None, None]: + """Create real limiters on a configuration and close them after the test.""" + created: list[QuotaLimiter] = [] + + def _make( + config: QuotaHandlersConfiguration, limiter: QuotaLimiterConfiguration + ) -> QuotaLimiter: + """Create the limiter object a configuration entry describes.""" + cls = ( + UserQuotaLimiter if limiter.type == "user_limiter" else ClusterQuotaLimiter + ) + instance = cls(config, limiter.initial_quota, limiter.quota_increase) + created.append(instance) + return instance + + yield _make + for instance in created: + instance.connection.close() + + +@pytest.fixture(name="scheduler_connection") +def scheduler_connection_fixture() -> Generator[Callable[..., Any], None, None]: + """Open the scheduler's own connection to a configuration's database.""" + connections: list[Any] = [] + + def _open(config: QuotaHandlersConfiguration) -> Any: + """Connect the way the scheduler does.""" + connection = connect(config) + connections.append(connection) + return connection + + yield _open + for connection in connections: + connection.close() + + +# ========================================== +# Renewal statements +# ========================================== + +RESET = limiter_config( + "user_limiter", initial_quota=100, quota_increase=0, period=ELAPSED +) +INCREASE = limiter_config( + "user_limiter", initial_quota=0, quota_increase=50, period=ELAPSED +) +CLUSTER_RESET = limiter_config( + "cluster_limiter", initial_quota=1000, quota_increase=0, period=ELAPSED +) + + +def with_period( + limiter: QuotaLimiterConfiguration, period: str +) -> QuotaLimiterConfiguration: + """Return the same limiter with another period.""" + return limiter.model_copy(update={"period": period}) + + +@pytest.mark.parametrize( + ("limiter", "renewed_balance"), + [ + pytest.param(RESET, 100, id="reset_to_initial_quota"), + pytest.param(INCREASE, -30 + 50, id="increase_by_quota_increase"), + ], +) +@pytest.mark.parametrize( + ("period", "elapsed"), + [ + pytest.param(ELAPSED, True, id="period_elapsed"), + pytest.param(NOT_ELAPSED, False, id="period_not_elapsed"), + ], +) +def test_renewal_applies_only_once_the_period_elapsed( # pylint: disable=too-many-arguments,too-many-positional-arguments + limiter: QuotaLimiterConfiguration, + renewed_balance: int, + period: str, + elapsed: bool, + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + scheduler_connection: Callable[..., Any], +) -> None: + """A consumed balance is renewed when, and only when, its period has elapsed.""" + config = handlers(quota_db, with_period(limiter, period)) + spend(make_limiter(config, config.limiters[0]), 30, ALICE) + consumed_balance = read_available(quota_db, ALICE, USER_SUBJECT) + assert consumed_balance == limiter.initial_quota - 30 + back_date(quota_db) + + renew(scheduler_connection(config), config) + + expected = renewed_balance if elapsed else consumed_balance + assert read_available(quota_db, ALICE, USER_SUBJECT) == expected + + +@pytest.mark.parametrize( + ("limiter", "renewed", "spent", "renewed_again"), + [ + pytest.param(RESET, 100, 90, 100, id="reset"), + pytest.param(INCREASE, 20, 10, 60, id="increase"), + ], +) +def test_renewal_happens_once_per_period( # pylint: disable=too-many-arguments,too-many-positional-arguments + limiter: QuotaLimiterConfiguration, + renewed: int, + spent: int, + renewed_again: int, + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + scheduler_connection: Callable[..., Any], +) -> None: + """A renewal moves the stamp, so the next sync waits for a whole period again. + + Both statements must move it. Before LCORE-4606 a row written by the + scheduler itself was renewed on every tick, which turned any + ``quota_increase`` into an unlimited quota. The last renewal works on the + stamp the scheduler wrote, in SQLite's own ``datetime('now')`` text. The + period is ten seconds so that a slow runner can not let it elapse between + the first and the second renewal. + """ + config = handlers(quota_db, with_period(limiter, NOT_ELAPSED)) + alice = make_limiter(config, config.limiters[0]) + connection = scheduler_connection(config) + spend(alice, 30, ALICE) + back_date(quota_db, seconds=PAST_NOT_ELAPSED) + + renew(connection, config) + assert read_available(quota_db, ALICE, USER_SUBJECT) == renewed + stamped_at = read_revoked_at(quota_db, ALICE, USER_SUBJECT) + + spend(alice, 10, ALICE) + renew(connection, config) + assert ( + read_available(quota_db, ALICE, USER_SUBJECT) == spent + ), "must not renew twice in one period" + assert read_revoked_at(quota_db, ALICE, USER_SUBJECT) == stamped_at + + back_date(quota_db, seconds=PAST_NOT_ELAPSED, style="scheduler") + renew(connection, config) + assert read_available(quota_db, ALICE, USER_SUBJECT) == renewed_again + + +@pytest.mark.parametrize( + ("limiter", "row_id", "subject"), + [ + pytest.param(RESET, ALICE, USER_SUBJECT, id="user"), + pytest.param(CLUSTER_RESET, CLUSTER, CLUSTER_SUBJECT, id="cluster"), + ], +) +def test_the_limiter_stamps_a_new_row_with_the_current_time_as_sqlite_reads_it( # pylint: disable=too-many-arguments,too-many-positional-arguments + limiter: QuotaLimiterConfiguration, + row_id: str, + subject: str, + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + scheduler_connection: Callable[..., Any], +) -> None: + """A row's period starts when the limiter creates it, in a format SQLite parses. + + This is the limiter's half of LCORE-4606: the renewal statements read + ``revoked_at`` through ``datetime()``, which yields NULL (and never + matches) for a missing stamp or one in a text it does not understand. + The stamp is checked as SQLite computes it, not as Python would. + """ + config = handlers(quota_db, with_period(limiter, NOT_ELAPSED)) + spend(make_limiter(config, config.limiters[0]), 30, row_id) + consumed_balance = limiter.initial_quota - 30 + connection = scheduler_connection(config) + + age = read_row_age_seconds(quota_db, row_id, subject) + assert age is not None, "SQLite can not read the stamp the limiter wrote" + assert 0 <= age < 5, f"the row should have been stamped just now, not {age}s ago" + + renew(connection, config) + assert read_available(quota_db, row_id, subject) == consumed_balance + + back_date(quota_db, seconds=PAST_NOT_ELAPSED) + renew(connection, config) + assert read_available(quota_db, row_id, subject) == limiter.initial_quota + + +@pytest.mark.parametrize( + ("limiter", "expected_calls"), + [ + pytest.param(RESET, {"reset_quota"}, id="zero_increase_skips_increase"), + pytest.param(INCREASE, {"increase_quota"}, id="zero_initial_skips_reset"), + ], +) +def test_the_statement_of_a_zero_value_is_not_executed( # pylint: disable=too-many-arguments,too-many-positional-arguments + limiter: QuotaLimiterConfiguration, + expected_calls: set[str], + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + scheduler_connection: Callable[..., Any], + mocker: MockerFixture, +) -> None: + """A zero increase or a zero initial quota runs no statement at all (LCORE-4606). + + Running the statement anyway stamped ``revoked_at`` without changing the + balance, which kept the other statement from ever matching the row. The + balance shows that for a zero increase (the reset after it then applies, + see ``test_renewal_applies_only_once_the_period_elapsed``) but not for a + zero initial quota, since the reset runs second; hence the spies. + """ + config = handlers(quota_db, limiter) + spend(make_limiter(config, limiter), 30, ALICE) + back_date(quota_db) + spies = { + name: mocker.spy(quota_scheduler, name) + for name in ("increase_quota", "reset_quota") + } + + renew(scheduler_connection(config), config) + + assert {name for name, spy in spies.items() if spy.called} == expected_calls + + +@pytest.mark.xfail( + strict=True, + reason=( + "the increase runs first and stamps revoked_at, so the reset of the " + "same tick no longer matches: the balance becomes 70 + 50 = 120 " + "instead of the 100 + 50 the quota_limiter module documents; " + "reported with LCORE-4610, to be fixed under its own ticket" + ), +) +def test_initial_quota_and_increase_together_reset_to_their_sum( + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + scheduler_connection: Callable[..., Any], +) -> None: + """With both values set, an elapsed period resets to ``initial_quota + quota_increase``.""" + both = limiter_config( + "user_limiter", initial_quota=100, quota_increase=50, period=ELAPSED + ) + config = handlers(quota_db, both) + spend(make_limiter(config, both), 30, ALICE) + back_date(quota_db) + + renew(scheduler_connection(config), config) + + assert read_available(quota_db, ALICE, USER_SUBJECT) == 100 + 50 + + +def test_renewal_touches_only_the_rows_of_its_own_subject( + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + scheduler_connection: Callable[..., Any], +) -> None: + """A cluster renewal leaves user rows alone and vice versa; all rows of a subject renew.""" + config = handlers(quota_db, RESET, CLUSTER_RESET) + users = make_limiter(config, RESET) + spend(users, 30, ALICE) + spend(users, 40, BOB) + spend(make_limiter(config, CLUSTER_RESET), 500) + back_date(quota_db) + connection = scheduler_connection(config) + + renew(connection, handlers(quota_db, CLUSTER_RESET)) + assert read_available(quota_db, CLUSTER, CLUSTER_SUBJECT) == 1000 + assert read_available(quota_db, ALICE, USER_SUBJECT) == 70 + assert read_available(quota_db, BOB, USER_SUBJECT) == 60 + + renew(connection, handlers(quota_db, RESET)) + assert read_available(quota_db, ALICE, USER_SUBJECT) == 100 + assert read_available(quota_db, BOB, USER_SUBJECT) == 100 + assert read_available(quota_db, CLUSTER, CLUSTER_SUBJECT) == 1000 + + +# ========================================== +# The scheduler loop +# ========================================== + + +def test_scheduler_loop_renews_and_then_sleeps_for_its_period( + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + mocker: MockerFixture, +) -> None: + """One pass of ``quota_scheduler``: connect, create the table, renew, sleep.""" + config = handlers(quota_db, RESET, scheduler_period=5) + spend(make_limiter(config, RESET), 30, ALICE) + back_date(quota_db) + sleep = mocker.patch("runners.quota_scheduler.sleep", side_effect=StopLoop) + + with pytest.raises(StopLoop): + quota_scheduler.quota_scheduler(config) + + assert read_available(quota_db, ALICE, USER_SUBJECT) == 100 + sleep.assert_called_once_with(5) + + +def test_scheduler_loop_creates_the_quota_table_on_an_empty_database( + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + mocker: MockerFixture, +) -> None: + """On a database no limiter has opened yet, the scheduler creates the table itself. + + In the deployed service the scheduler thread starts before the first + request builds the limiters, so on a fresh deployment it is the one that + must create ``quota_limits``; its first ticks would otherwise fail on a + missing table, silently, forever. + """ + config = handlers(quota_db, RESET) + init_tables = mocker.spy(quota_scheduler, "init_tables") + mocker.patch("runners.quota_scheduler.sleep", side_effect=StopLoop) + assert not quota_db.exists() + + with pytest.raises(StopLoop): + quota_scheduler.quota_scheduler(config) + + assert table_exists(quota_db, "quota_limits") + init_tables.assert_called_once_with(mocker.ANY, CREATE_QUOTA_TABLE_SQLITE) + spend(make_limiter(config, RESET), 30, ALICE) + assert read_available(quota_db, ALICE, USER_SUBJECT) == 70 + + +def test_scheduler_loop_survives_a_failing_limiter_and_renews_the_next( + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + mocker: MockerFixture, +) -> None: + """A database error on one limiter is logged; the others are still renewed. + + This is what keeps the scheduler thread alive after a transient error, + such as SQLite's ``database is locked``. + """ + config = handlers(quota_db, RESET, CLUSTER_RESET) + spend(make_limiter(config, RESET), 30, ALICE) + spend(make_limiter(config, CLUSTER_RESET), 500) + back_date(quota_db) + real_revocation = quota_scheduler.quota_revocation + revocation = mocker.patch("runners.quota_scheduler.quota_revocation") + + def failing_once(*args: Any, **kwargs: Any) -> None: + """Fail the first limiter of the tick, renew the others for real.""" + if revocation.call_count == 1: + raise sqlite3.OperationalError("database is locked") + real_revocation(*args, **kwargs) + + revocation.side_effect = failing_once + sleep = mocker.patch("runners.quota_scheduler.sleep", side_effect=StopLoop) + + with pytest.raises(StopLoop): + quota_scheduler.quota_scheduler(config) + + assert revocation.call_count == len(config.limiters) + assert read_available(quota_db, ALICE, USER_SUBJECT) == 70, "the failed one" + assert read_available(quota_db, CLUSTER, CLUSTER_SUBJECT) == 1000 + sleep.assert_called_once_with(config.scheduler.period) + + +def test_scheduler_loop_reconnects_when_its_connection_died( + quota_db: Path, + make_limiter: Callable[..., QuotaLimiter], + mocker: MockerFixture, +) -> None: + """A connection lost between two ticks is replaced and the renewal goes on.""" + config = handlers(quota_db, RESET) + alice = make_limiter(config, RESET) + spend(alice, 30, ALICE) + back_date(quota_db) + connect_spy = mocker.spy(quota_scheduler, "connect") + + def between_ticks(_period: int) -> None: + """Kill the scheduler's connection after the first tick, stop after the second.""" + if connect_spy.call_count == 1: + assert read_available(quota_db, ALICE, USER_SUBJECT) == 100 + connect_spy.spy_return.close() + spend(alice, 10, ALICE) + back_date(quota_db) + return + raise StopLoop + + mocker.patch("runners.quota_scheduler.sleep", side_effect=between_ticks) + + with pytest.raises(StopLoop): + quota_scheduler.quota_scheduler(config) + + assert connect_spy.call_count == 2 + assert read_available(quota_db, ALICE, USER_SUBJECT) == 100 From e3661b6fa4604e18f9f654541bc614157717f579 Mon Sep 17 00:00:00 2001 From: Maxim Svistunov Date: Mon, 2 Nov 2026 13:32:38 +0100 Subject: [PATCH 3/3] LCORE-4610: test how the endpoints check and charge quota Add integration tests that send requests through the real endpoints with real limiters on SQLite, so the client's view and the database are compared. They cover /v1/query, /v1/streaming_query, /v1/responses (blocking and streamed) and /v1/infer: a turn is charged and the balance reported, an exhausted user is rejected with the 429 body, a shield-blocked streamed turn costs nothing, and /v1/infer is not enforced without a quota subject. They also cover the shared cluster balance, which limiter is named when both are exhausted, the token history, a limiter reconnecting after its connection died, and an unreachable database giving a 500 instead of a served turn, also when the error comes after the model answered. --- tests/integration/endpoints/README.md | 4 + .../endpoints/test_quota_endpoints.py | 874 ++++++++++++++++++ 2 files changed, 878 insertions(+) create mode 100644 tests/integration/endpoints/test_quota_endpoints.py diff --git a/tests/integration/endpoints/README.md b/tests/integration/endpoints/README.md index 479f6179d..9d0cd3e6e 100644 --- a/tests/integration/endpoints/README.md +++ b/tests/integration/endpoints/README.md @@ -76,6 +76,10 @@ Integration tests for the /query endpoint BYOK inline and tool RAG functionality Integration tests for the /query endpoint (using Responses API). +## [test_quota_endpoints.py](test_quota_endpoints.py) + +Integration tests for quota enforcement on the endpoints (LCORE-4610). + ## [test_responses_byok_integration.py](test_responses_byok_integration.py) Integration tests for the /responses endpoint BYOK RAG functionality. diff --git a/tests/integration/endpoints/test_quota_endpoints.py b/tests/integration/endpoints/test_quota_endpoints.py new file mode 100644 index 000000000..40b26f1a7 --- /dev/null +++ b/tests/integration/endpoints/test_quota_endpoints.py @@ -0,0 +1,874 @@ +"""Integration tests for quota enforcement on the endpoints (LCORE-4610). + +Every endpoint that calls the model checks the caller's quota first and +charges the usage of the call afterwards: ``/v1/query``, +``/v1/streaming_query``, ``/v1/responses`` (blocking and streaming) and +``/v1/infer`` when ``rlsapi_v1.quota_subject`` is set. The first three also +report the balances left, keyed by limiter class. + +The quota is kept by real limiters on a real SQLite file and the handlers are +the real ones; only the model is mocked, and it reports the same usage for +every call. The balances are asserted both as the client sees them and as the +database holds them. +""" + +# pylint: disable=too-many-arguments,too-many-positional-arguments +# pylint: disable=protected-access,redefined-outer-name + +import json +import sqlite3 +from collections.abc import AsyncIterator, Callable, Generator +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +import pytest +from fastapi import BackgroundTasks, HTTPException, Request, status +from fastapi.responses import StreamingResponse +from fastapi.testclient import TestClient +from ogx_client.models.open_ai_response_object_stream_response_completed import ( + OpenAIResponseObjectStreamResponseCompleted, +) +from ogx_client.models.open_ai_response_object_stream_response_created import ( + OpenAIResponseObjectStreamResponseCreated, +) +from ogx_client.models.open_ai_response_object_stream_response_in_progress import ( + OpenAIResponseObjectStreamResponseInProgress, +) +from pytest_mock import AsyncMockType, MockerFixture +from sqlalchemy.orm import Session + +from app.endpoints.query import query_endpoint_handler +from app.endpoints.responses import responses_endpoint_handler +from app.endpoints.rlsapi_v1 import infer_endpoint +from app.endpoints.streaming_query import streaming_query_endpoint_handler +from authentication.interface import AuthTuple +from authentication.rh_identity import RHIdentityData +from configuration import AppConfig +from models.api.requests import QueryRequest, ResponsesRequest +from models.api.requests.rlsapi import RlsapiV1InferRequest +from models.api.responses.successful import ResponsesResponse +from models.common.moderation import ShieldModerationBlocked +from models.common.responses.contexts import ResponsesContext +from models.config import QuotaLimiterConfiguration +from models.database.conversations import UserConversation +from quota.quota_limiter import QuotaLimiter +from tests.integration._quota_helpers import ( + CLUSTER, + CLUSTER_SUBJECT, + USER_SUBJECT, + handlers, + limiter_config, + read_available, + read_token_usage, +) +from tests.integration.conftest import ( + TEST_MODEL, + TEST_MODEL_NAME, + TEST_PROVIDER, + make_openai_response_object, + set_query_agent_run, + set_streaming_query_agent_run, +) +from tests.unit.utils.auth_helpers import mock_authorization_resolvers + +TURN_INPUT, TURN_OUTPUT = 100, 50 +TURN = TURN_INPUT + TURN_OUTPUT +"""What the model reports for every call it answers.""" + +QUESTION = "What is Ansible?" +ALICE = "11111111-1111-1111-1111-111111111111" +BOB = "22222222-2222-2222-2222-222222222222" +ORG = "org-42" +SYSTEM = "cn-7" + +QUOTA_EXCEEDED = "The quota has been exceeded" +DATABASE_ERROR = { + "response": "Database query failed", + "cause": "Failed to query the database", +} + + +def user_quota(initial_quota: int) -> QuotaLimiterConfiguration: + """A user limiter that is never renewed during a test.""" + return limiter_config( + "user_limiter", initial_quota=initial_quota, quota_increase=0, period="1 day" + ) + + +def cluster_quota(initial_quota: int) -> QuotaLimiterConfiguration: + """A cluster limiter that is never renewed during a test.""" + return limiter_config( + "cluster_limiter", + initial_quota=initial_quota, + quota_increase=0, + period="1 day", + ) + + +def no_tokens_for_user(user_id: str) -> dict[str, str]: + """The 429 body a user with an exhausted quota gets.""" + return { + "response": QUOTA_EXCEEDED, + "cause": f"User {user_id} has no available tokens", + } + + +NO_TOKENS_FOR_CLUSTER = { + "response": QUOTA_EXCEEDED, + "cause": "Cluster has no available tokens", +} + + +def auth_for(user_id: str) -> AuthTuple: + """The auth tuple the noop authentication produces for a user.""" + return (user_id, "lightspeed-user", True, "") + + +# ========================================== +# Fixtures +# ========================================== + + +@dataclass +class Quota: + """The quota of a test: its configuration and its database. + + Attributes: + config: The configuration, with the quota handlers set. + db_path: The SQLite file the limiters and the history write to. + + """ + + config: AppConfig + db_path: Path + + @property + def limiters(self) -> list[QuotaLimiter]: + """The real limiters built from the configuration.""" + return self.config.quota_limiters + + def user_balance(self, user_id: str) -> Optional[int]: + """What the database holds for a user, None when it never saw the user.""" + return read_available(self.db_path, user_id, USER_SUBJECT) + + def cluster_balance(self) -> Optional[int]: + """What the database holds for the cluster, None when no row exists.""" + return read_available(self.db_path, CLUSTER, CLUSTER_SUBJECT) + + +ConfigureQuota = Callable[..., Quota] + + +@pytest.fixture(name="quota") +def quota_fixture( + test_config: AppConfig, tmp_path: Path +) -> Generator[ConfigureQuota, None, None]: + """Let a test choose its limiters; they live on a SQLite file in ``tmp_path``. + + Yields: + A function taking the limiter configurations (and ``token_history``) + that installs them in the loaded configuration and returns the + :class:`Quota` of the test. + + """ + assert test_config._configuration is not None + db_path = tmp_path / "quota.db" + + def _configure( + *limiters: QuotaLimiterConfiguration, token_history: bool = False + ) -> Quota: + """Install the given limiters and reset what was built before.""" + assert test_config._configuration is not None + test_config._configuration.quota_handlers = handlers( + db_path, *limiters, enable_token_history=token_history + ) + test_config._quota_limiters = [] + test_config._token_usage_history = None + return Quota(test_config, db_path) + + yield _configure + for limiter in test_config._quota_limiters: + try: + limiter.connection.close() + except sqlite3.ProgrammingError: + # opened by the thread TestClient served a request from; SQLite + # refuses to close it from here and closes it with the object + pass + test_config._quota_limiters = [] + history = test_config._token_usage_history + if history is not None and history.connection is not None: + history.connection.close() + test_config._token_usage_history = None + + +async def _responses_stream(response_object: Any) -> AsyncIterator[Any]: + """Stream a response the way OGX does: created, in progress, completed.""" + yield OpenAIResponseObjectStreamResponseCreated( + response=response_object, sequence_number=0, type="response.created" + ) + yield OpenAIResponseObjectStreamResponseInProgress( + response=response_object, sequence_number=1, type="response.in_progress" + ) + yield OpenAIResponseObjectStreamResponseCompleted( + response=response_object, sequence_number=2, type="response.completed" + ) + + +@pytest.fixture(name="model") +def model_fixture( + test_config: AppConfig, + mock_ogx_client: AsyncMockType, + mock_query_agent: AsyncMockType, + mock_streaming_query_agent: AsyncMockType, + mocker: MockerFixture, +) -> AsyncMockType: + """Make the model report the same usage to every endpoint. + + The query endpoints answer through the agent, ``/v1/responses`` and + ``/v1/infer`` through ``responses.create`` of the OGX client. The + responses context is built without validating the mocked client, as in + the other responses integration tests. + + Returns: + The mocked OGX client. + + """ + set_query_agent_run( + mock_query_agent, mocker, input_tokens=TURN_INPUT, output_tokens=TURN_OUTPUT + ) + set_streaming_query_agent_run( + mock_streaming_query_agent, + mocker, + input_tokens=TURN_INPUT, + output_tokens=TURN_OUTPUT, + ) + + async def _create(**kwargs: Any) -> Any: + """Answer like OGX: a response object, or a stream of events.""" + answer = make_openai_response_object( + input_tokens=TURN_INPUT, output_tokens=TURN_OUTPUT + ) + return _responses_stream(answer) if kwargs.get("stream") else answer + + mock_ogx_client.responses.create = mocker.AsyncMock(side_effect=_create) + mocker.patch( + "app.endpoints.responses.ResponsesContext", + side_effect=ResponsesContext.model_construct, + ) + + rlsapi_holder = mocker.patch("app.endpoints.rlsapi_v1.AsyncOgxClientHolder") + rlsapi_holder.return_value.get_client.return_value = mock_ogx_client + test_config.inference.default_model = TEST_MODEL_NAME + test_config.inference.default_provider = TEST_PROVIDER + mock_authorization_resolvers(mocker) + return mock_ogx_client + + +# ========================================== +# Requests +# ========================================== + + +async def query(request: Request, auth: AuthTuple) -> Any: + """Send a query as the given user.""" + return await query_endpoint_handler( + request=request, + query_request=QueryRequest(query=QUESTION), + auth=auth, + mcp_headers={}, + ) + + +async def _events(response: Any) -> list[dict[str, Any]]: + """Read a streaming response to its end and parse its ``data:`` lines.""" + assert isinstance(response, StreamingResponse) + events: list[dict[str, Any]] = [] + async for chunk in response.body_iterator: + for line in str(chunk).splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + events.append(json.loads(line[len("data: ") :])) + return events + + +async def streaming_query(request: Request, auth: AuthTuple) -> dict[str, Any]: + """Send a streaming query as the given user and return its ``end`` event.""" + response = await streaming_query_endpoint_handler( + request=request, + query_request=QueryRequest(query=QUESTION), + auth=auth, + mcp_headers={}, + ) + (end,) = [event for event in await _events(response) if event["event"] == "end"] + return end + + +async def responses(request: Request, auth: AuthTuple, stream: bool) -> Any: + """Send a responses request as the given user. + + Returns: + The response object when blocking, else the ``(type, available_quotas)`` + of every streamed event that carries a response object. + + """ + response = await responses_endpoint_handler( + request=request, + responses_request=ResponsesRequest( + input=QUESTION, + model=TEST_MODEL, + stream=stream, + generate_topic_summary=False, + ), + auth=auth, + mcp_headers={}, + ) + if not stream: + assert isinstance(response, ResponsesResponse) + return response + return [ + (event["type"], event["response"].get("available_quotas")) + for event in await _events(response) + if "response" in event + ] + + +async def infer(request: Request, auth: AuthTuple) -> Any: + """Send an rlsapi v1 inference as the given user.""" + return await infer_endpoint( + infer_request=RlsapiV1InferRequest(question=QUESTION), + request=request, + background_tasks=BackgroundTasks(), + auth=auth, + ) + + +# ========================================== +# /v1/query +# ========================================== + + +@pytest.mark.asyncio +async def test_query_charges_the_turn_and_reports_the_balance( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, +) -> None: + """A turn costs its input and output tokens; the client and the database agree.""" + _ = model + scene = quota(user_quota(1000)) + + response = await query(test_request, test_auth) + + assert (response.input_tokens, response.output_tokens) == (TURN_INPUT, TURN_OUTPUT) + assert response.available_quotas == {"UserQuotaLimiter": 1000 - TURN} + assert scene.user_balance(test_auth[0]) == 1000 - TURN + + +@pytest.mark.asyncio +async def test_query_is_served_while_the_balance_is_positive_then_rejected( + quota: ConfigureQuota, + model: AsyncMockType, + mock_query_agent: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, + patch_db_session: Session, +) -> None: + """One token lets one turn through; the next is a 429 that costs nothing. + + The rejected request calls no model, creates no conversation and leaves + the balance where the served turn put it. + """ + _ = model + user_id = test_auth[0] + scene = quota(user_quota(1)) + + served = await query(test_request, test_auth) + assert served.available_quotas == {"UserQuotaLimiter": 1 - TURN} + + with pytest.raises(HTTPException) as rejected: + await query(test_request, test_auth) + + assert rejected.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS + assert rejected.value.detail == no_tokens_for_user(user_id) + assert mock_query_agent.run.await_count == 1 + assert patch_db_session.query(UserConversation).count() == 1 + assert scene.user_balance(user_id) == 1 - TURN + + +def test_query_over_http_returns_the_429_body( + integration_http_client: TestClient, + quota: ConfigureQuota, + model: AsyncMockType, +) -> None: + """The quota error reaches an HTTP client as a 429 with the documented body.""" + _ = model + scene = quota(user_quota(1)) + url = f"/v1/query?user_id={ALICE}" + + served = integration_http_client.post(url, json={"query": QUESTION}) + assert served.status_code == status.HTTP_200_OK, served.text + assert served.json()["available_quotas"] == {"UserQuotaLimiter": 1 - TURN} + + rejected = integration_http_client.post(url, json={"query": QUESTION}) + assert rejected.status_code == status.HTTP_429_TOO_MANY_REQUESTS + assert rejected.json() == {"detail": no_tokens_for_user(ALICE)} + assert scene.user_balance(ALICE) == 1 - TURN + + +# ========================================== +# /v1/streaming_query +# ========================================== + + +@pytest.mark.asyncio +async def test_streaming_query_reports_the_balance_in_its_end_event( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, +) -> None: + """The stream's end event carries the usage of the turn and the balance after it.""" + _ = model + scene = quota(user_quota(1000)) + + end = await streaming_query(test_request, test_auth) + + assert end["data"]["input_tokens"] == TURN_INPUT + assert end["data"]["output_tokens"] == TURN_OUTPUT + assert end["available_quotas"] == {"UserQuotaLimiter": 1000 - TURN} + assert scene.user_balance(test_auth[0]) == 1000 - TURN + + +@pytest.mark.asyncio +async def test_streaming_query_of_an_exhausted_user_is_rejected_before_streaming( + quota: ConfigureQuota, + model: AsyncMockType, + mock_streaming_query_agent: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, +) -> None: + """An exhausted user gets the 429 as a plain error, never a stream.""" + _ = model + quota(user_quota(1)) + await streaming_query(test_request, test_auth) + + with pytest.raises(HTTPException) as rejected: + await streaming_query(test_request, test_auth) + + assert rejected.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS + assert rejected.value.detail == no_tokens_for_user(test_auth[0]) + assert mock_streaming_query_agent.run_stream_events.call_count == 1 + + +# ========================================== +# /v1/responses +# ========================================== + + +@pytest.mark.asyncio +async def test_blocking_responses_reports_the_balance( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, +) -> None: + """A blocking response carries the balance left after its usage was charged.""" + _ = model + scene = quota(user_quota(1000)) + + response = await responses(test_request, test_auth, stream=False) + + assert response.available_quotas == {"UserQuotaLimiter": 1000 - TURN} + assert scene.user_balance(test_auth[0]) == 1000 - TURN + + +@pytest.mark.asyncio +async def test_streamed_responses_reports_the_balance_on_the_terminal_event( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, +) -> None: + """Intermediate events carry no balance; the completed event carries the charged one.""" + _ = model + scene = quota(user_quota(1000)) + + events = await responses(test_request, test_auth, stream=True) + + assert events == [ + ("response.created", None), + ("response.in_progress", {}), + ("response.completed", {"UserQuotaLimiter": 1000 - TURN}), + ] + assert scene.user_balance(test_auth[0]) == 1000 - TURN + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["blocking", "streaming"]) +async def test_responses_of_an_exhausted_user_is_rejected( + stream: bool, + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, +) -> None: + """Blocking or streaming, the second request of a one-token user is a 429.""" + quota(user_quota(1)) + await responses(test_request, test_auth, stream) + + with pytest.raises(HTTPException) as rejected: + await responses(test_request, test_auth, stream) + + assert rejected.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS + assert rejected.value.detail == no_tokens_for_user(test_auth[0]) + assert model.responses.create.await_count == 1 + + +@pytest.mark.asyncio +async def test_a_shield_blocked_streamed_responses_turn_costs_nothing( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, + mocker: MockerFixture, +) -> None: + """A turn an input shield blocks calls no model and is not charged. + + Its stream still ends with the balance, the one the served turn before + it left behind. + """ + scene = quota(user_quota(1000)) + await responses(test_request, test_auth, stream=True) + balance = 1000 - TURN + assert scene.user_balance(test_auth[0]) == balance + mocker.patch( + "app.endpoints.responses.run_shield_moderation_v2", + return_value=ShieldModerationBlocked( + message="Content blocked by safety shield", moderation_id="modr-blocked" + ), + ) + + events = await responses(test_request, test_auth, stream=True) + + assert events == [ + ("response.created", {}), + ("response.completed", {"UserQuotaLimiter": balance}), + ] + assert model.responses.create.await_count == 1 + assert scene.user_balance(test_auth[0]) == balance + + +# ========================================== +# /v1/infer +# ========================================== + + +def _user_identity(org_id: str) -> RHIdentityData: + """The rh-identity data of a console user of an organization.""" + return RHIdentityData( + { + "identity": { + "type": "User", + "org_id": org_id, + "user": {"user_id": "console-user", "username": "console@example"}, + } + } + ) + + +def _system_identity(cn: str) -> RHIdentityData: + """The rh-identity data of a certificate-authenticated system.""" + return RHIdentityData( + {"identity": {"type": "System", "org_id": ORG, "system": {"cn": cn}}} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("quota_subject", "identity", "charged"), + [ + pytest.param("user_id", None, ALICE, id="user_id"), + pytest.param("org_id", _user_identity(ORG), ORG, id="org_id"), + pytest.param("system_id", _system_identity(SYSTEM), SYSTEM, id="system_id"), + pytest.param("org_id", None, ALICE, id="org_id_without_identity_falls_back"), + pytest.param( + "system_id", None, ALICE, id="system_id_without_identity_falls_back" + ), + ], +) +async def test_infer_charges_the_configured_quota_subject( + quota_subject: str, + identity: Optional[RHIdentityData], + charged: str, + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, +) -> None: + """The row charged is the one ``rlsapi_v1.quota_subject`` names, or the user's.""" + _ = model + scene = quota(user_quota(1000)) + scene.config.configuration.rlsapi_v1.quota_subject = quota_subject # type: ignore[assignment] + if identity is not None: + test_request.state.rh_identity_data = identity + + await infer(test_request, auth_for(ALICE)) + + assert scene.user_balance(charged) == 1000 - TURN + for other in {ALICE, ORG, SYSTEM} - {charged}: + assert scene.user_balance(other) is None + + +@pytest.mark.asyncio +async def test_infer_of_an_exhausted_user_is_rejected( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, +) -> None: + """With a quota subject set, a one-token user gets one inference, then a 429.""" + scene = quota(user_quota(1)) + scene.config.configuration.rlsapi_v1.quota_subject = "user_id" + await infer(test_request, auth_for(ALICE)) + + with pytest.raises(HTTPException) as rejected: + await infer(test_request, auth_for(ALICE)) + + assert rejected.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS + assert rejected.value.detail == no_tokens_for_user(ALICE) + assert model.responses.create.await_count == 1 + assert scene.user_balance(ALICE) == 1 - TURN + + +@pytest.mark.asyncio +async def test_infer_is_not_enforced_without_a_quota_subject( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, +) -> None: + """Without ``quota_subject`` the limiters are neither consulted nor charged.""" + scene = quota(user_quota(1)) + assert scene.config.configuration.rlsapi_v1.quota_subject is None + + await infer(test_request, auth_for(ALICE)) + await infer(test_request, auth_for(ALICE)) + + assert model.responses.create.await_count == 2 + assert not scene.db_path.exists(), "no limiter was built, so no database either" + + +# ========================================== +# User and cluster limiters together +# ========================================== + + +@pytest.mark.asyncio +async def test_the_cluster_balance_is_shared_and_both_balances_are_reported( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, +) -> None: + """Each user has a balance of their own; the cluster balance is paid by all of them.""" + _ = model + scene = quota(user_quota(1000), cluster_quota(1000)) + + alice = await query(test_request, auth_for(ALICE)) + bob = await query(test_request, auth_for(BOB)) + + assert alice.available_quotas == { + "UserQuotaLimiter": 1000 - TURN, + "ClusterQuotaLimiter": 1000 - TURN, + } + assert bob.available_quotas == { + "UserQuotaLimiter": 1000 - TURN, + "ClusterQuotaLimiter": 1000 - 2 * TURN, + } + assert scene.user_balance(ALICE) == 1000 - TURN + assert scene.user_balance(BOB) == 1000 - TURN + assert scene.cluster_balance() == 1000 - 2 * TURN + + +@pytest.mark.asyncio +async def test_an_exhausted_cluster_rejects_a_user_who_still_has_tokens( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, +) -> None: + """The cause names the cluster; the user's own balance is created but not charged.""" + _ = model + scene = quota(user_quota(1000), cluster_quota(1)) + await query(test_request, auth_for(ALICE)) + assert scene.cluster_balance() == 1 - TURN + + with pytest.raises(HTTPException) as rejected: + await query(test_request, auth_for(BOB)) + + assert rejected.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS + assert rejected.value.detail == NO_TOKENS_FOR_CLUSTER + assert scene.user_balance(BOB) == 1000 + assert scene.cluster_balance() == 1 - TURN + + +@pytest.mark.asyncio +async def test_an_exhausted_user_is_named_before_the_cluster_is_consulted( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, +) -> None: + """Limiters are checked in configuration order; the user limiter comes first.""" + _ = model + scene = quota(user_quota(1), cluster_quota(1000)) + await query(test_request, auth_for(ALICE)) + + with pytest.raises(HTTPException) as rejected: + await query(test_request, auth_for(ALICE)) + + assert rejected.value.detail == no_tokens_for_user(ALICE) + assert scene.cluster_balance() == 1000 - TURN + + +# ========================================== +# Token usage history +# ========================================== + + +@pytest.mark.asyncio +@pytest.mark.parametrize("with_limiter", [True, False], ids=["limiter", "no_limiter"]) +async def test_token_history_accumulates_per_user_provider_and_model( + with_limiter: bool, + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, +) -> None: + """With ``enable_token_history`` every turn is added to its user's row. + + The history is kept with or without limiters: the storage alone is enough. + """ + _ = model + limiters = [user_quota(1000)] if with_limiter else [] + scene = quota(*limiters, token_history=True) + + await query(test_request, auth_for(ALICE)) + await query(test_request, auth_for(ALICE)) + bob = await query(test_request, auth_for(BOB)) + + assert read_token_usage(scene.db_path) == { + (ALICE, TEST_PROVIDER, TEST_MODEL_NAME): (2 * TURN_INPUT, 2 * TURN_OUTPUT), + (BOB, TEST_PROVIDER, TEST_MODEL_NAME): (TURN_INPUT, TURN_OUTPUT), + } + expected_quotas = {"UserQuotaLimiter": 1000 - TURN} if with_limiter else {} + assert bob.available_quotas == expected_quotas + + +# ========================================== +# The quota database goes away +# ========================================== + + +@pytest.mark.asyncio +async def test_a_limiter_reconnects_after_its_connection_died( + quota: ConfigureQuota, + model: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, +) -> None: + """A lost connection is reopened on the next request, which continues the balance.""" + _ = model + scene = quota(user_quota(1000)) + await query(test_request, test_auth) + (limiter,) = scene.limiters + limiter.connection.close() + + response = await query(test_request, test_auth) + + assert response.available_quotas == {"UserQuotaLimiter": 1000 - 2 * TURN} + assert scene.user_balance(test_auth[0]) == 1000 - 2 * TURN + + +@pytest.mark.asyncio +async def test_an_unreachable_quota_database_is_a_500_not_a_served_turn( + quota: ConfigureQuota, + model: AsyncMockType, + mock_query_agent: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, + mocker: MockerFixture, +) -> None: + """When the database cannot be reopened the check fails closed with a database error.""" + _ = model + scene = quota(user_quota(1000)) + await query(test_request, test_auth) + (limiter,) = scene.limiters + limiter.connection.close() + mocker.patch( + "quota.quota_limiter.connect_sqlite", + side_effect=sqlite3.OperationalError("unable to open database file"), + ) + + with pytest.raises(HTTPException) as failed: + await query(test_request, test_auth) + + assert failed.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert failed.value.detail == DATABASE_ERROR + assert mock_query_agent.run.await_count == 1 + + +def _fail_on_the_second_call(mocker: MockerFixture, limiter: QuotaLimiter) -> None: + """Let the balance read of the check pass, then fail the one of the report.""" + real_available_quota = limiter.available_quota + + def read_then_fail(subject_id: str = "") -> int: + """Delegate once, then raise like a lost database.""" + if available_quota.call_count > 1: + raise sqlite3.OperationalError("database is locked") + return real_available_quota(subject_id) + + available_quota = mocker.patch.object( + limiter, "available_quota", side_effect=read_then_fail + ) + + +def _fail_the_charge(mocker: MockerFixture, limiter: QuotaLimiter) -> None: + """Fail the consumption of the turn's usage.""" + mocker.patch.object( + limiter, + "consume_tokens", + side_effect=sqlite3.OperationalError("database is locked"), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("fail", "balance_left"), + [ + pytest.param(_fail_the_charge, 1000, id="charging_the_turn"), + pytest.param(_fail_on_the_second_call, 1000 - TURN, id="reporting_the_balance"), + ], +) +async def test_a_database_error_after_the_model_answered_is_a_500( + fail: Callable[[MockerFixture, QuotaLimiter], None], + balance_left: int, + quota: ConfigureQuota, + model: AsyncMockType, + mock_query_agent: AsyncMockType, + test_request: Request, + test_auth: AuthTuple, + mocker: MockerFixture, +) -> None: + """The check passed and the model answered; then the database fails. + + Whether the usage can not be charged or the balance can not be read for + the response, the client gets the database error. A charge that went + through before the read failed stays charged. + """ + _ = model + scene = quota(user_quota(1000)) + (limiter,) = scene.limiters + fail(mocker, limiter) + + with pytest.raises(HTTPException) as failed: + await query(test_request, test_auth) + + assert failed.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + assert failed.value.detail == DATABASE_ERROR + assert mock_query_agent.run.await_count == 1 + assert scene.user_balance(test_auth[0]) == balance_left