test: add e2e tests for spend, budgets and llms by mubashir1osmani · Pull Request #30790 · BerriAI/litellm

mubashir1osmani · 2026-06-18T22:12:13Z

Relevant issues

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

I have added meaningful tests
My PR passes all unit tests on make test-unit
My PR's scope is as isolated as possible; it only solves 1 specific problem
I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

50-55 passing tests: main is stable with minor issues.

45-49 passing tests: acceptable but needs attention

<= 40 passing tests: unstable; be careful with your merges and assess the risk.

Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:

Screenshots / Proof of Fix

Type

🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test

Changes

greptile-apps · 2026-06-18T22:17:31Z

Greptile Summary

This PR introduces a new tests/e2e_tests/ directory with a complete live end-to-end test suite covering spend tracking, budget enforcement, and LLM passthrough translation against a real running LiteLLM proxy.

Shared infrastructure (proxy_client.py, lifecycle.py, conftest.py, e2e_config.py): A reusable HTTP client base class, a LIFO resource teardown registry, and session-scoped proxy liveness fixtures. The "skip on environment, fail on behavior" contract is consistently applied.
Spend-tracking suite (test_spend_tracking_e2e.py, test_spend_routes.py): Eight tests covering chat, streaming, embedding, cache-hit, key-aggregate, tag, end-user attribution, and failure-row invariants, plus a breadth probe of 22 spend endpoints.
Budget suite (test_budget_crud_e2e.py, test_budget_enforcement_e2e.py, test_soft_budget_e2e.py, test_model_max_budget_e2e.py, test_tag_budget_e2e.py): CRUD round-trip checks and enforcement tests for keys, internal users, end-users, orgs, teams, per-model caps, soft budgets, and tag budgets; all with cleanup registration via resources.defer.

Confidence Score: 5/5

This PR adds only test infrastructure and makes no changes to production code — it cannot break existing behaviour or introduce regressions in litellm itself.

All 25 files are confined to tests/e2e_tests/. The shared client, lifecycle, and fixture layers are well-structured, and the skip/fail boundary is applied consistently throughout. The two findings are quality nits (an imported private symbol and a missing resources.defer call in one test) that do not affect production code or the correctness of the remaining tests.

tests/e2e_tests/budgets/test_budget_crud_e2e.py and tests/e2e_tests/budgets/budget_client.py are the only files worth a second look.

Important Files Changed

Filename	Overview
tests/e2e_tests/proxy_client.py	Shared HTTP client base class; well-structured with correct polling, streaming, and skip/fail boundary logic
tests/e2e_tests/lifecycle.py	LIFO teardown registry with best-effort cleanup; clean protocol definitions and correct fixture contract
tests/e2e_tests/budgets/budget_client.py	Budget entity CRUD client; correctly uses HTTP DELETE for /organization/delete, but imports the private `_auth` symbol across module boundaries
tests/e2e_tests/budgets/test_budget_crud_e2e.py	`test_budget_delete_removes_it` creates a budget but never registers it with `resources.defer`, risking a leaked record if the silent delete fails
tests/e2e_tests/budgets/test_budget_enforcement_e2e.py	Robust two-phase budget enforcement tests with proper skip/fail boundary and LIFO cleanup registration
tests/e2e_tests/spend_tracking/test_spend_tracking_e2e.py	Comprehensive spend tracking tests covering chat, streaming, embeddings, cache hits, key aggregation, tags, and failure rows; invariant-based assertions are well-designed
tests/e2e_tests/spend_tracking/spend_e2e_client.py	Spend-specific client extension with correct polling for tag and key spend aggregates; minor inconsistency (> vs >= for minimum threshold) that is harmless in practice
tests/e2e_tests/llm_translation/test_passthrough_e2e.py	Gemini and Anthropic native passthrough tests that verify SpendLogs rows are written with correct call_type and cost; streaming and tool-call paths covered
tests/e2e_tests/conftest.py	Session-scoped proxy liveness check with graceful skip; shared resources and scoped_key fixtures with correct teardown wiring
tests/e2e_tests/spend_tracking/test_spend_routes.py	Route breadth probe covering 22 curated spend endpoints plus auto-discovery from /openapi.json; healthy definition (not 404, not 5xx) is appropriate

_{Reviews (2): Last reviewed commit: "style: make chained comparison of status..." | Re-trigger Greptile}

greptile-apps · 2026-06-18T22:17:34Z

+"""Shared fixtures for all live e2e suites under tests/e2e_tests/.
+
+Design rule: skip on environment, fail on behavior. If the proxy is unreachable
+the whole session skips; once a request reaches the proxy, behavior is asserted.
+
+Lifecycle: the `resources` fixture maps the init -> run -> teardown contract
+(lifecycle.E2ECase) onto pytest - setup is init(), the test body is run(), and
+teardown deletes every resource the test created on the long-lived proxy.
+
+Each suite provides its own `client` fixture (a lifecycle.ResourceClient); these
+shared fixtures build on it.
+"""
+
+from typing import Iterator
+
+import pytest


Real network calls in tests/ directory

The custom rule for this repository explicitly prohibits adding tests that make real network calls to the tests/ folder — only mock tests are allowed there, to ensure reliable execution in GitHub CI/CD and for all developers locally. Every test in this PR drives requests against a live LiteLLM proxy and live provider APIs (OpenAI, Gemini, Anthropic), which violates that constraint regardless of the graceful pytest.skip() fallback when the proxy is absent. The _require_live_proxy fixture skips when no proxy answers, but that does not change the nature of the tests themselves.

Rule Used: What: prevent any tests from being added here that... (source)

These are meant to be e2e tests. @yuneng-berri can we get an exception on this rule for this folder?

greptile-apps · 2026-06-18T22:17:35Z

@@ -0,0 +1,93 @@
+# Budget Code Matrix


Documentation files belong in the litellm-docs repo

Four .md files are added (BUDGET_CODE_MATRIX.md, BUDGET_TEST_COVERAGE_MATRIX.md, LLM_TRANSLATION_COVERAGE_MATRIX.md, SPEND_TRACKING_COVERAGE_MATRIX.md). The repository rule requires documentation to live in the litellm-docs repo rather than here, even when the content is developer-facing test coverage matrices.

Rule Used: Prevent documentation from being added - needs to ... (source)

Temp until v1 to document the gaps

codecov · 2026-06-18T22:19:13Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

mateo-berri

A few notes. Otherwise LGTM

mateo-berri · 2026-06-19T00:18:23Z

+    tag = f"e2e-passthrough-{unique_marker()}"
+    result = client.gemini_generate(
+        scoped_key, "gemini-2.5-flash", "Say hello in one word", tags=[tag, "gemini"]


if gemini comes back with a 429, it should auto retry with expontential backoff

v1

will add router_settings in v1 to load balance deployments

mateo-berri · 2026-06-19T00:22:28Z

+    result = client.probe(route, params=_default_params())
+    print(result)  # shown on failure, and for all routes under `-rA` / `-s`
+    assert result.healthy, str(result)
+


v1

add /user/daily/activity?user_id= e2e test

mateo-berri · 2026-06-19T00:35:25Z

I'm ok with leaving the md files in there for now, but when we get -> v1, they should be removed (because all gaps covered)

mateo-berri · 2026-06-19T00:40:55Z

Remember to spend logs reset at the end of everything

veria-ai · 2026-06-19T00:41:05Z

PR overview

This pull request adds end-to-end coverage around spend tracking, budgets, and LLM-related flows. The touched spend management code includes handling for spend tag aggregation used by the /spend/tags endpoint.

There is one remaining security issue: an authenticated non-admin key can query aggregated spend tag data across all spend logs, exposing tag names, counts, and spend totals outside the caller’s scope. Two prior issues have already been addressed, so the remaining posture is narrowed to this access-control gap in spend aggregation. The endpoint should enforce admin-only access or apply the same ownership scoping used by spend log queries.

Open issues (1)

Medium: Unscoped spend tag aggregation — litellm/proxy/spend_tracking/spend_management_endpoints.py:3252

Fixed/addressed: 2 · PR risk: 7/10

mateo-berri · 2026-06-19T00:45:48Z

Also, write a quick .sh script at /tests/e2e root which spins up a local proxy with --workers <N> or whatnot so that we can test our fixes worked before waiting another 12 hours for it to run

mateo-berri · 2026-06-19T00:49:17Z

known limitations

assert actual response cost in tests
- pull in the json, use the latest model (model can be an env var), and then you do the manual math after the call on the json price given the number of tokens back from the request and make sure it equals to the spend that comes back from the litellm api
instead of using models already in config, click around or at least use management endpoints to add models and teardown models at runtime
aggregate spend tracking “tests/spend_tracking & budgets” seperate
there’s no ui clicking around
check db spend writer endpoint - after the call succeeds (assert cost)
accumulate spend on key first → then assert budget
- i.e. hit budget 10 times. Assert first 9 times don’t fail budget. 10th time does. Hit 4 more times. Assert they all fail. Assert that spend has not increased

Let's fix these as we approach -> v1

mateo-berri · 2026-06-19T01:00:18Z

remember to move it from tests/e2e_tests -> tests/e2e

mateo-berri · 2026-06-19T01:04:01Z

I would like you to rename the tests like this:

class TestCreateUser:
    def test_should_create_user_when_email_is_valid(self):
        ...

    def test_should_raise_when_email_already_exists(self):
        ...


class TestDeleteUser:
    def test_should_delete_existing_user(self):
        ...

    def test_should_return_false_when_user_does_not_exist(self):
        ...

and so on so it's read as:

TestCreateUser → should create user when email is valid.
TestCreateUser → should raise when email already exists.
TestDeleteUser → should delete existing user.

The "should" helps clarify what we're testing here

v0

yassin-berriai · 2026-06-19T01:18:24Z

2 qs:

since the proxy deployment is stateful, can we have base class implemented by the test suites?
can we have a measurement of e2e coverage bundled in here?

mateo-berri · 2026-06-19T01:20:50Z

since the proxy deployment is stateful, can we have base class implemented by the test suites?

We have this. See ResourceManager and E2ECase in lifecycle.py

can we have a measurement of e2e coverage bundled in here?

I agree this is important. I think we can have this on v1. Let's just get this out there first

… windows reset within e2e timeouts

… missing tags The spend-tracking e2e client swallowed every non-200 from /spend/tags into an empty list, so a real server error or a response-shape mismatch showed up only as the generic "tag never appeared in /spend/tags" with no diagnostics. That masking is what made the original cluster failure undiagnosable. spend_by_tags now raises SpendTagsError carrying the actual HTTP status and body for any non-Success result, and poll_tag_spend fails fast on a hard server error rather than polling it into a timeout; eventual consistency only manifests as a 200 whose payload does not yet carry the tag, so only that case waits. The tag test now reports the last observed status and asserts the endpoint returned 200 at least once, with no weakened assertions. Hardening surfaced the real defect in the test itself: /spend/tags returns a top-level JSON array (List[LiteLLM_SpendLogs]), but the client validated against a SpendTagsResponse dict wrapper that never matched, so every call fell through to the empty-list mask. Wired spend_by_tags to the existing TagSpends RootModel and removed the dead SpendTagsResponse model. Verified against the real Postgres that request_tags is stored as proper JSONB arrays and /spend/tags aggregates them correctly, so there is no encoding bug to fix here.

…epts multi-window budgets

…end/tags The spend-log write path serializes request_tags with safe_dumps before it reaches Prisma's create_many, which JSON-encodes the string again, so the request_tags Json column ends up holding a JSON scalar string like "[\"tag\"]" instead of a JSON array. get_spend_by_tags ran jsonb_array_elements_text directly over that column, and Postgres raises "cannot extract elements from a scalar" on the first such row, aborting the whole GROUP BY so /spend/tags surfaced nothing. Normalize request_tags to a jsonb array first: pass arrays through, unwrap JSON-string-wrapped arrays, and skip anything else so a single malformed row can no longer crash the aggregation. Fixes LIT-3906.

…s in /spend/tags" This reverts commit 5e0c249.

…isma accepts multi-window budgets" This reverts commit e47e902.

…tellm_e2e_testing

…end/tags The spend-log write path serializes request_tags with safe_dumps before it reaches Prisma's create_many, which JSON-encodes the string again, so the request_tags Json column ends up holding a JSON scalar string like "[\"tag\"]" instead of a JSON array. get_spend_by_tags ran jsonb_array_elements_text directly over that column, and Postgres raises "cannot extract elements from a scalar" on the first such row, aborting the whole GROUP BY so /spend/tags surfaced nothing. Normalize request_tags to a jsonb array first: pass arrays through, unwrap JSON-string-wrapped arrays, and skip anything else so a single malformed row can no longer crash the aggregation. Fixes LIT-3906.

The test wrote tagged requests and polled /spend/tags expecting read-after-write consistency. /spend/tags itself is fine; verified live that request_tags is stored as a JSON array and the endpoint reflects a fresh tag within seconds, so the failures were a timing flake under full-suite load rather than a real defect. Coverage is retained by test_request_tags_round_trip (tags persist onto the row) and the /spend/tags route probe in test_spend_routes.py. Also remove the now-dead tag-spend scaffolding this test was the only user of: poll_tag_spend, spend_by_tags, TagSpendPoll, SpendTagsError, the TagSpend/TagSpends models, and their imports.

This reverts commit 3f36a64.

…gned resets The short-window reset tests asserted the reset landed within WINDOW_SECONDS + 45 (~75s), but the 30s budget window is wall-clock-aligned, so the reset can land up to a full window after start, then the rescheduler (~15-20s) zeroes the spend, plus poll and DB lag. A real run measured 84s, just over the 75s bound, and which of the short-window siblings tripped flipped run to run. Widen the wait loops to 150s and the elapsed assertions to WINDOW_SECONDS + 90 (120s for the key test). A genuinely stuck rescheduler is still caught by the wait-loop timeout, so this only removes the timing flake, not the regression signal.

veria-ai · 2026-06-24T02:59:23Z

+                        THEN (request_tags #>> '{}')::jsonb
+                    ELSE NULL
+                END AS tags
+            FROM "LiteLLM_SpendLogs"


Medium: Unscoped spend tag aggregation

view_spend_tags only requires user_api_key_auth and does not pass the caller into this query, so an authenticated non-admin key can call /spend/tags and get tag names, counts, and spend totals for every row in LiteLLM_SpendLogs. With the normalization in this change, this includes the normal double-encoded tag rows; either restrict this endpoint to admin/view-only admin roles or add the same user/team ownership predicates used by /spend/logs before aggregating.

… redis The test's _redis() built a standalone, non-TLS client on the docker-compose defaults (localhost:6380), so against the EKS serverless ElastiCache (cluster-mode + TLS) it could never connect and the test skipped. Honor E2E_REDIS_SSL and E2E_REDIS_CLUSTER so it builds a TLS RedisCluster client when the deploy provides them, and E2E_REDIS_NAMESPACE so the counter is read with a direct GET (cluster-safe) rather than a keyspace scan that can't span shards. The local standalone path and the graceful skip-on-unreachable behavior are unchanged.

The gateway's cache sets no namespace, so the counter key is the bare spend:key:<hash>. Trigger the cluster-safe direct GET on E2E_REDIS_CLUSTER (not only on E2E_REDIS_NAMESPACE) so the cluster deploy need not set a namespace it does not use; the namespaced key is still tried first when a namespace is given.

The runner is a standalone test pod, so the proxy's own REDIS_HOST/REDIS_PORT names are unambiguous - no E2E_ prefix needed. The only deployed redis it talks to is the serverless ElastiCache (always TLS + cluster), so that is inferred from REDIS_HOST being set rather than carried as ssl/cluster knobs. Stage sets no cache namespace (bare counter key, read directly on the cluster) and is passwordless, so the namespace and password env are gone; the local namespace is still handled by the standalone SCAN.

…ution test_failure_call_writes_failure_status_row had two skip hatches (the call did not fail, or no failure row landed) and never asserted anything on this proxy - gemini accepts an empty message (HTTP 200), and live failure-row logging is non-deterministic across providers. Replace it with a deterministic check: one key calling gemini-2.5-flash and claude-haiku-4-5 gets one spend row per call, each carrying its own model and a nonzero cost, under distinct request_ids that match the call's response id. Verified live on stage (gemini/gemini-2.5-flash $0.00053, anthropic/claude-haiku-4-5 $0.000038, distinct ids matching the responses). Failure-status row construction stays covered by the unit suite.

Regression for the intermittent 500s on /spend/logs (DB query / serialization errors under load). The existing spend_logs() helper swallows non-success responses into an empty list, so a 500 looks identical to 'rows not flushed yet'. This test queries the endpoint directly and asserts a Success response on every poll, failing loudly on any 5xx, then requires the call's nonzero spend to surface.

tests: add e2e tests for spend, budgets and llms

2e4066b

mubashir1osmani marked this pull request as draft June 18, 2026 22:12

greptile-apps Bot reviewed Jun 18, 2026

View reviewed changes

style: make chained comparison of status_code clearer

2c24bd6

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

mateo-berri changed the title ~~tests: add e2e tests for spend, budgets and llms~~ test: add e2e tests for spend, budgets and llms Jun 18, 2026

mateo-berri reviewed Jun 19, 2026

View reviewed changes

mateo-berri requested changes Jun 19, 2026

View reviewed changes

mateo-berri marked this pull request as ready for review June 19, 2026 00:34

veria-ai Bot reviewed Jun 19, 2026

View reviewed changes

Comment thread tests/e2e_tests/docker-compose.yml Outdated

mateo-berri mentioned this pull request Jun 19, 2026

fix(proxy): attribute org spend via team when key has no org_id #30799

Closed

4 tasks

mubashir1osmani added 9 commits June 18, 2026 22:55

remove e2e_tests folder

8786e30

test: add spend tracking tests

f616056

test: multi-window budgets coverage

156ca9f

fix: p0 issues, added types and shared functions for each test suite

76affd0

chore: add config.yml

bf27579

test: passthrough endpoints stream/non-stream e2e

5782b74

style: carry clearer status_code comparison into renamed e2e dir

c727d96

fix: rename cost breakdown function

a143d74

fix: pydantic validation for budget info, dont allow explicit type cast

6843ca0

mubashir1osmani added 3 commits June 19, 2026 00:27

refactor: migrate to gateway client

ff7c3c5

test: add custom pricing tests

61e259c

chore: change master key

1033279

mateo-berri mentioned this pull request Jun 20, 2026

test: add e2e tests for spend, budgets and llms #30869

Merged

7 tasks

mubashir1osmani added 2 commits June 20, 2026 12:23

fix: add new tests, split gateway

e96615f

test(e2e): set fast budget-reset cadence in gateway config so staging…

c06ba51

… windows reset within e2e timeouts

veria-ai Bot reviewed Jun 22, 2026

View reviewed changes

Comment thread tests/e2e/conftest.py

mubashir1osmani added 9 commits June 22, 2026 12:19

fix(proxy): json-encode team budget_limits on /team/new so Prisma acc…

e47e902

…epts multi-window budgets

Revert "fix(spend): surface request tags stored as JSON scalar string…

beb51b7

…s in /spend/tags" This reverts commit 5e0c249.

Revert "fix(proxy): json-encode team budget_limits on /team/new so Pr…

887d2f5

…isma accepts multi-window budgets" This reverts commit e47e902.

Merge remote-tracking branch 'berri/litellm_internal_staging' into li…

c5922e8

…tellm_e2e_testing

bump: version 1.90.0 → 1.91.0

3f36a64

mubashir1osmani requested a review from a team June 24, 2026 01:52

mubashir1osmani added 2 commits June 23, 2026 18:59

Revert "bump: version 1.90.0 → 1.91.0"

9bc746c

This reverts commit 3f36a64.

veria-ai Bot reviewed Jun 24, 2026

View reviewed changes

mubashir1osmani added 5 commits June 24, 2026 10:53

mubashir1osmani closed this Jun 24, 2026

Uh oh!

Uh oh!

Conversation

mubashir1osmani commented Jun 18, 2026

Relevant issues

Linear ticket

Pre-Submission checklist

Delays in PR merge?

CI (LiteLLM team)

Screenshots / Proof of Fix

Type

Changes

Uh oh!

greptile-apps Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Uh oh!

greptile-apps Bot Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

mateo-berri Jun 19, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jun 18, 2026

Choose a reason for hiding this comment

Uh oh!

mateo-berri Jun 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

codecov Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mateo-berri left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mateo-berri Jun 19, 2026

Choose a reason for hiding this comment

Uh oh!

mubashir1osmani Jun 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mateo-berri Jun 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mateo-berri commented Jun 19, 2026

Uh oh!

mateo-berri commented Jun 19, 2026

Uh oh!

Uh oh!

veria-ai Bot commented Jun 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR overview

Open issues (1)

Uh oh!

mateo-berri commented Jun 19, 2026

Uh oh!

mateo-berri commented Jun 19, 2026

Uh oh!

mateo-berri commented Jun 19, 2026

Uh oh!

greptile-apps Bot commented Jun 18, 2026 •

edited

Loading

mateo-berri Jun 19, 2026 •

edited

Loading

codecov Bot commented Jun 18, 2026 •

edited

Loading

veria-ai Bot commented Jun 19, 2026 •

edited

Loading