Repository navigation
LCORE-4610: unit and integration tests for quota management - #3134
Open
max-svistunov wants to merge 3 commits into
Open
max-svistunov wants to merge 3 commits into
max-svistunov wants to merge 3 commits into
Conversation
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.
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.
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.
6 of 20 tasks
Contributor
|
Warning Review limit reachedYou've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Next included review available in 31 minutes. View limit details
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. CodeRabbit Verify
|
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Part 1 of 2 for LCORE-4610: unit and integration tests for quota management. Tests only; nothing under
src/changes. Part 2 (#3135) adds the e2e feature on top of this.Integration (real limiters on a real SQLite file, real handlers, only the model mocked):
tests/integration/test_quota_scheduler.py: the periodic renewal that nothing executed so far. Reset and increase apply only once a limiter's period has elapsed and exactly once per period (both statements move therevoked_atstamp; a row the scheduler itself stamped renews again), the stamp the limiter writes on a new row is read back through SQLite's owndatetime()and is the current time (the limiter's half of LCORE-4606: a stamp SQLite can not parse is never renewed, silently), a zero increase or zero initial quota runs no statement, a renewal touches only its own subject, one pass of the scheduler loop renews and sleeps for its period, the scheduler createsquota_limitson a database no limiter has opened (on a fresh deployment it starts before the first request builds the limiters), a database error on one limiter is logged and the next limiter is still renewed, and a connection lost between two ticks is replaced. Time is never waited for:revoked_atis moved into the past, with a ten second period wherever a slow runner could otherwise let a two second one elapse between two steps.tests/integration/endpoints/test_quota_endpoints.py:/v1/query,/v1/streaming_query(end event),/v1/responsesblocking and streaming (in_progresscarries{}, the terminal event the balance), and/v1/inferfor everyquota_subject(user id, org id, system id, and the fallbacks) charge the turn and report the balance; a one-token user is served once and then gets the documented 429 body, also over HTTP, calling no model and creating no conversation; a streamed/v1/responsesturn an input shield blocks calls no model, costs nothing and still ends with the balance; user and cluster limiters together, with the first exhausted one naming the cause; token history accumulation with and without limiters; a limiter reconnects after its connection died; an unreachable database is a 500 before the model is called, and a database that fails after the model answered (charging the turn, or reading the balance for the response) is a 500 as well.tests/unit/runners/test_quota_scheduler.py(statement selection, subject mapping, early returns, reconnection attempts, PostgreSQL statement shape with a mocked connection; the order of the increase and the reset is deliberately not pinned, see below) andtests/unit/quota/test_token_usage_history.py.PostgreSQL gets shape tests only: no PostgreSQL service exists in CI or locally.
One behaviour is documented rather than asserted: with
initial_quotaandquota_increaseboth set, thequota.quota_limitermodule promises a reset to their sum (150 for 100 + 50). The scheduler runs the two statements separately and each stampsrevoked_at, so whichever runs first hides the other: the code gives 120 today, and swapping the order would give 100, never 150.test_initial_quota_and_increase_together_reset_to_their_sumisxfail(strict=True)with that reason, so the discrepancy is visible and the test turns red the day it is fixed. Ticket to be filed.Type of change
pyproject.toml+uv.lock]requirements.*.txtfor Konflux]Tools used to create PR
Identify any AI code assistants used in this PR (for transparency and review context)
Related Tickets & Documents
Checklist before requesting a review
Testing
Run the new tests:
Expected: 26 unit tests pass; 41 integration tests pass and one is reported as XFAIL (
test_initial_quota_and_increase_together_reset_to_their_sum, see the description).What was run here: the new tests pass and were run three times in a row without a flake; the full
tests/unitandtests/integrationsuites (withoutcontainer_lifecycle) pass; ruff, black, pylint and pyright report nothing new on the touched files. The tests were also checked against hand-made production mutants (an increase that does not move the stamp, a new row with a NULL stamp, a scheduler that skips table creation, a revocation error that escapes the loop, among others): every one of them is caught by at least one of the new tests, except the increase/reset order, which is left unpinned on purpose.