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/endpoints/README.md b/tests/integration/endpoints/README.md index 7eb541917..553a21c1f 100644 --- a/tests/integration/endpoints/README.md +++ b/tests/integration/endpoints/README.md @@ -80,6 +80,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 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 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