scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test(integration): add spend log alias probe cells for the daily activity routes

Browse source
This commit is contained in:
mateo-berri 2026-09-29 20:05:50 -07:00
parent 59d6c462d4
commit 54da4e523b
3 changed files with 602 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

65
tests/integration/_support/daily_activity.py
View file

@ -3,6 +3,8 @@ import uuid
from collections.abc import Iterator, Mapping, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timedelta
from hashlib import sha256
from itertools import chain
from typing import Final
@ -34,7 +36,16 @@ INSERT_SPEND_LOG: Final = (
" VALUES (%s, 'acompletion', %s, %s::timestamp, %s::timestamp, %s)"
)
DELETE_SPEND_LOG: Final = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
INSERT_SPEND_LOG_ROW: Final = (
'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, "startTime", "endTime", metadata, team_id, "user")'
" VALUES (%s, 'acompletion', %s, %s::timestamp, %s::timestamp, %s, %s, %s)"
)
DELETE_SPEND_LOG_ROWS: Final = 'DELETE FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)'
DELETE_KEY_ROW: Final = 'DELETE FROM "LiteLLM_VerificationToken" WHERE token = %s'
DELETE_ARCHIVED_KEY_ROW: Final = 'DELETE FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s'
LOCK_TABLE: Final = sql.SQL("LOCK TABLE {table} IN ACCESS EXCLUSIVE MODE")
SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs"
FIRST_SPEND_LOG_AT: Final = datetime(2026, 2, 3, 12, 0, 0)
@dataclass(frozen=True, slots=True)
@ -67,6 +78,10 @@ def key_no_key_table_holds() -> str:
return f"integration-ownerless-{uuid.uuid4().hex}"
def digest_no_key_table_holds() -> str:
return sha256(uuid.uuid4().bytes).hexdigest()
def activity_of_key(
gateway: Gateway, path: str, api_key: str, *, reader: str | None = None, **filters: str
) -> httpx.Response:
@ -148,6 +163,56 @@ def spend_log_naming_only_an_alias(request_id: str, api_key: str, started: str,
connection.execute(DELETE_SPEND_LOG, (request_id,))
@dataclass(frozen=True, slots=True)
class SpendLogRow:
started: str
metadata: JsonValue = None
team_id: str | None = None
user: str | None = None
def started_at(index: int) -> str:
return (FIRST_SPEND_LOG_AT + timedelta(seconds=index)).strftime("%Y-%m-%d %H:%M:%S")
def nameless_rows(count: int, first_index: int = 0) -> tuple[SpendLogRow, ...]:
return tuple(SpendLogRow(started_at(first_index + offset), {}) for offset in range(count))
def named_row(index: int, alias: str) -> SpendLogRow:
return SpendLogRow(started_at(index), {"user_api_key_alias": alias})
@contextmanager
def spend_logs_of_key(
api_key: str, rows: Sequence[SpendLogRow], *, database_url: str | None = None
) -> Iterator[tuple[str, ...]]:
request_ids: Final = tuple(f"integration-{uuid.uuid4().hex}" for _ in rows)
with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection:
connection.cursor().executemany(
INSERT_SPEND_LOG_ROW,
tuple(
(request_id, api_key, row.started, row.started, Jsonb(row.metadata), row.team_id, row.user)
for request_id, row in zip(request_ids, rows, strict=True)
),
)
try:
yield request_ids
finally:
delete_spend_logs(request_ids, database_url=database_url)
def delete_spend_logs(request_ids: Sequence[str], *, database_url: str | None = None) -> None:
with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection:
connection.execute(DELETE_SPEND_LOG_ROWS, (list(request_ids),))
def purge_key_from_the_key_tables(digest: str, *, database_url: str | None = None) -> None:
with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection:
connection.execute(DELETE_KEY_ROW, (digest,))
connection.execute(DELETE_ARCHIVED_KEY_ROW, (digest,))
@contextmanager
def locked_table(table: str, *, database_url: str | None = None) -> Iterator[None]:
with psycopg.connect(database_url or os.environ["DATABASE_URL"]) as connection:

490
tests/integration/spend/test_daily_activity_key_alias_probes.py Normal file
View file

@ -0,0 +1,490 @@
import time
import uuid
from collections.abc import Callable, Iterator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value
from integration._support.daily_activity import (
AGGREGATED_USER_ACTIVITY,
DAY,
ROUTES,
SPEND_LOGS_TABLE,
USER_SPEND,
Route,
SpendLogRow,
activity_of_key,
assert_key_reported,
daily_rows,
digest_no_key_table_holds,
key_metadata,
locked_table,
named_row,
nameless_rows,
records_of_key,
seeded_metrics,
seeded_row,
spend_logs_of_key,
started_at,
user_row,
user_with_an_email,
)
from integration._support.database import read_rows, scratch_database
from integration._support.process import OwnedProxy, owned_proxy_process
from pydantic import JsonValue
DAY_OUTSIDE_THE_WINDOW: Final = "2026-02-10"
GIVES_UP_WITHIN_SECONDS: Final = 10
CONCURRENT_READS: Final = 20
CACHED_MISS_CLEARS_WITHIN_SECONDS: Final = 45
ALIAS_OF_ONE_SPEND_LOG: Final = (
"SELECT metadata->>'user_api_key_alias' AS alias FROM \"LiteLLM_SpendLogs\" WHERE request_id = %s"
)
def _alias() -> str:
return f"integration-alias-{uuid.uuid4().hex}"
def _named_between_fifty_and_fifty(alias: str) -> tuple[SpendLogRow, ...]:
return (*nameless_rows(50), named_row(50, alias), *nameless_rows(50, 51))
def _oldest_named(alias: str) -> tuple[SpendLogRow, ...]:
return (named_row(0, alias), *nameless_rows(150, 1))
def _newest_named(alias: str) -> tuple[SpendLogRow, ...]:
return (*nameless_rows(150), named_row(150, alias))
def _both_edges_named(alias: str) -> tuple[SpendLogRow, ...]:
return (named_row(0, alias), *nameless_rows(150, 1), named_row(151, alias))
def _named_after_one_hundred(alias: str) -> tuple[SpendLogRow, ...]:
return (*nameless_rows(100), named_row(100, alias), *nameless_rows(99, 101))
def _named_after_ninety_nine(alias: str) -> tuple[SpendLogRow, ...]:
return (*nameless_rows(99), named_row(99, alias), *nameless_rows(100, 100))
def _named_only_in_the_middle(alias: str) -> tuple[SpendLogRow, ...]:
return (*nameless_rows(100), named_row(100, alias), *nameless_rows(100, 101))
def _renamed_and_renamed_back(alias: str, other: str) -> tuple[SpendLogRow, ...]:
return (
named_row(0, alias),
*nameless_rows(100, 1),
named_row(101, other),
*nameless_rows(100, 102),
named_row(202, alias),
)
def _team_in_the_column(team: str) -> SpendLogRow:
return SpendLogRow(started_at(0), {}, team_id=team)
def _team_in_the_metadata(team: str) -> SpendLogRow:
return SpendLogRow(started_at(0), {"user_api_key_team_id": team})
def _user_in_the_column(user: str) -> SpendLogRow:
return SpendLogRow(started_at(0), {}, user=user)
def _user_in_the_metadata(user: str) -> SpendLogRow:
return SpendLogRow(started_at(0), {"user_api_key_user_id": user})
def _activity_on_route(gateway: Gateway, route: Route, api_key: str, entity: str) -> httpx.Response:
filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity}
return activity_of_key(gateway, route.path, api_key, **filters)
def _reported_aliases(response: httpx.Response, api_key: str) -> tuple[JsonValue, ...]:
if response.status_code != 200:
return ()
return tuple(
object_value(object_value(record)["metadata"])["key_alias"]
for record in records_of_key(object_value(response.json()), api_key)
)
def _names_the_key(api_key: str, alias: str) -> Callable[[httpx.Response], bool]:
def names(response: httpx.Response) -> bool:
reported: Final = _reported_aliases(response, api_key)
return bool(reported) and frozenset(reported) == frozenset((alias,))
return names
@contextmanager
def _proxy_on(gateway: Gateway, directory: Path, database_url: str, *, workers: int = 1) -> Iterator[OwnedProxy]:
with owned_proxy_process(
gateway,
directory,
{"DATABASE_URL": database_url},
remove_environment=("DATABASE_URL_READ_REPLICA",),
workers=workers,
) as owned:
yield owned
def _owner_on(candidate: Gateway) -> tuple[str, str]:
owner: Final = f"integration-{uuid.uuid4().hex}"
email: Final = f"{owner}@example.com"
candidate.post("/user/new", {"user_id": owner, "user_email": email, "auto_create_key": False})
return owner, email
@pytest.mark.parametrize("route", ROUTES, ids=lambda route: route.path.strip("/").replace("/", "_"))
def test_alias_named_only_by_a_spend_log_is_reported_on_every_daily_activity_route(
gateway: Gateway, route: Route
) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
entity: Final = f"integration-entity-{uuid.uuid4().hex}"
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
entity_rows: Final = (
() if route.table == USER_SPEND else (seeded_row(route.table, route.entity_column, entity, api_key, DAY),)
)
filters: Final = {} if route.entity_filter is None else {route.entity_filter: entity}
with (
daily_rows((user_row(owner, api_key, DAY), *entity_rows)),
spend_logs_of_key(api_key, (named_row(0, alias),)),
):
assert_key_reported(
activity_of_key(gateway, route.path, api_key, **filters),
api_key,
DAY,
key_metadata(alias=alias, user=owner, email=email),
seeded_metrics(1),
)
@pytest.mark.parametrize(
"layout",
(
pytest.param(_named_between_fifty_and_fifty, id="named_between_50_and_50_nameless"),
pytest.param(_oldest_named, id="oldest_named_150_nameless_newer"),
pytest.param(_newest_named, id="newest_named_150_nameless_older"),
pytest.param(_both_edges_named, id="both_edges_named_150_nameless_between"),
pytest.param(_named_after_one_hundred, id="100_nameless_named_99_nameless"),
pytest.param(_named_after_ninety_nine, id="99_nameless_named_100_nameless"),
),
)
def test_alias_on_an_edge_of_the_window_is_reported_whatever_surrounds_it(
gateway: Gateway, layout: Callable[[str], tuple[SpendLogRow, ...]]
) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, layout(alias)):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(alias=alias, user=owner, email=email),
seeded_metrics(1),
)
def test_alias_named_only_in_the_middle_of_two_hundred_nameless_rows_is_not_picked_up(gateway: Gateway) -> None:
api_key: Final = digest_no_key_table_holds()
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
with (
daily_rows((user_row(owner, api_key, DAY),)),
spend_logs_of_key(api_key, _named_only_in_the_middle(_alias())),
):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(user=owner, email=email),
seeded_metrics(1),
)
def test_key_renamed_and_renamed_back_is_reported_with_the_alias_on_both_edges(gateway: Gateway) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
rows: Final = _renamed_and_renamed_back(alias, _alias())
with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(alias=alias, user=owner, email=email),
seeded_metrics(1),
)
@pytest.mark.parametrize(
"spend_log_of_team",
(
pytest.param(_team_in_the_column, id="team_id_column"),
pytest.param(_team_in_the_metadata, id="team_id_in_metadata"),
),
)
def test_team_named_only_by_a_spend_log_is_reported_next_to_the_daily_owner(
gateway: Gateway, spend_log_of_team: Callable[[str], SpendLogRow]
) -> None:
api_key: Final = digest_no_key_table_holds()
team: Final = f"integration-team-{uuid.uuid4().hex}"
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (spend_log_of_team(team),)):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(team=team, user=owner, email=email),
seeded_metrics(1),
)
@pytest.mark.parametrize(
"spend_log_of_user",
(
pytest.param(_user_in_the_column, id="user_column"),
pytest.param(_user_in_the_metadata, id="user_id_in_metadata"),
),
)
def test_user_named_by_a_spend_log_beats_the_owner_the_daily_rows_name(
gateway: Gateway, spend_log_of_user: Callable[[str], SpendLogRow]
) -> None:
api_key: Final = digest_no_key_table_holds()
with gateway.scenario() as scenario:
daily_owner, _ = user_with_an_email(scenario)
log_user, log_email = user_with_an_email(scenario)
with (
daily_rows((user_row(daily_owner, api_key, DAY),)),
spend_logs_of_key(api_key, (spend_log_of_user(log_user),)),
):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(user=log_user, email=log_email),
seeded_metrics(1),
)
def test_hashed_jwt_digest_is_named_by_its_spend_log(gateway: Gateway) -> None:
api_key: Final = f"hashed-jwt-{sha256(uuid.uuid4().bytes).hexdigest()}"
alias: Final = _alias()
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (named_row(0, alias),)):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(alias=alias, user=owner, email=email),
seeded_metrics(1),
)
@pytest.mark.parametrize(
("started", "inside_the_window"),
(
pytest.param("2026-02-01 23:59:59", False, id="second_before_the_window"),
pytest.param("2026-02-02 00:00:00", True, id="first_second_of_the_window"),
pytest.param("2026-02-04 23:59:59", True, id="last_second_of_the_window"),
pytest.param("2026-02-05 00:00:00", False, id="first_second_after_the_window"),
),
)
def test_spend_log_names_the_key_only_from_one_day_before_to_two_days_after_the_read(
gateway: Gateway, started: str, inside_the_window: bool
) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
row: Final = SpendLogRow(started, {"user_api_key_alias": alias})
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (row,)):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(alias=alias if inside_the_window else None, user=owner, email=email),
seeded_metrics(1),
)
def test_two_aliases_on_the_two_edges_leave_the_key_unnamed(gateway: Gateway) -> None:
api_key: Final = digest_no_key_table_holds()
rows: Final = (named_row(0, _alias()), *nameless_rows(150, 1), named_row(151, _alias()))
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(user=owner, email=email),
seeded_metrics(1),
)
@pytest.mark.parametrize(
"unnamed_rows",
(
pytest.param((SpendLogRow(started_at(0), {"user_api_key_alias": ""}),), id="empty_string_alias"),
pytest.param(
(SpendLogRow(started_at(0), ["x"]), SpendLogRow(started_at(1), "x")), id="array_then_string_metadata"
),
),
)
def test_rows_without_a_usable_alias_do_not_hide_the_named_row_after_them(
gateway: Gateway, unnamed_rows: tuple[SpendLogRow, ...]
) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
rows: Final = (*unnamed_rows, named_row(len(unnamed_rows), alias))
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, rows):
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(alias=alias, user=owner, email=email),
seeded_metrics(1),
)
@pytest.mark.parametrize(
"stored_alias",
(
pytest.param(123, id="json_int"),
pytest.param(["a"], id="json_list"),
pytest.param("a" * 5000, id="five_kb_string"),
),
)
def test_alias_of_an_unexpected_shape_is_reported_as_postgres_renders_it(
gateway: Gateway, stored_alias: JsonValue
) -> None:
api_key: Final = digest_no_key_table_holds()
row: Final = SpendLogRow(started_at(0), {"user_api_key_alias": stored_alias})
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
with daily_rows((user_row(owner, api_key, DAY),)), spend_logs_of_key(api_key, (row,)) as request_ids:
rendered: Final = read_rows(ALIAS_OF_ONE_SPEND_LOG, (request_ids[0],))[0]["alias"]
assert isinstance(rendered, str) and rendered, rendered
assert_key_reported(
activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key),
api_key,
DAY,
key_metadata(alias=rendered, user=owner, email=email),
seeded_metrics(1),
)
@pytest.mark.timeout(300)
def test_alias_found_once_is_served_from_the_cache_for_the_same_window_only(gateway: Gateway, tmp_path: Path) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned:
owner, email = _owner_on(owned.gateway)
rows: Final = (user_row(owner, api_key, DAY), user_row(owner, api_key, DAY_OUTSIDE_THE_WINDOW))
with daily_rows(rows, database_url=database_url):
with spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url):
first: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key)
cached: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key)
other_window: Final = owned.gateway.request(
"GET",
AGGREGATED_USER_ACTIVITY,
params={"start_date": DAY_OUTSIDE_THE_WINDOW, "end_date": DAY_OUTSIDE_THE_WINDOW, "api_key": api_key},
)
named: Final = key_metadata(alias=alias, user=owner, email=email)
assert_key_reported(first, api_key, DAY, named, seeded_metrics(1))
assert_key_reported(cached, api_key, DAY, named, seeded_metrics(1))
assert_key_reported(
other_window, api_key, DAY_OUTSIDE_THE_WINDOW, key_metadata(user=owner, email=email), seeded_metrics(1)
)
@pytest.mark.timeout(300)
def test_alias_logged_after_a_cached_miss_shows_once_the_miss_expires(gateway: Gateway, tmp_path: Path) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url) as owned:
owner, email = _owner_on(owned.gateway)
with daily_rows((user_row(owner, api_key, DAY),), database_url=database_url):
missed: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key)
with spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url):
named: Final = eventually(
lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key),
_names_the_key(api_key, alias),
seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS,
)
assert_key_reported(missed, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1))
assert_key_reported(named, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1))
@pytest.mark.timeout(300)
def test_alias_lookup_gives_up_while_spend_logs_are_locked_and_answers_once_they_are_not(
gateway: Gateway, tmp_path: Path
) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
with scratch_database() as database_url, _proxy_on(gateway, tmp_path, database_url, workers=2) as owned:
owner, email = _owner_on(owned.gateway)
with (
daily_rows((user_row(owner, api_key, DAY),), database_url=database_url),
spend_logs_of_key(api_key, (named_row(0, alias),), database_url=database_url),
):
with locked_table(SPEND_LOGS_TABLE, database_url=database_url):
started: Final = time.monotonic()
locked: Final = activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key)
waited: Final = time.monotonic() - started
unlocked: Final = eventually(
lambda: activity_of_key(owned.gateway, AGGREGATED_USER_ACTIVITY, api_key),
_names_the_key(api_key, alias),
seconds=CACHED_MISS_CLEARS_WITHIN_SECONDS,
)
assert waited < GIVES_UP_WITHIN_SECONDS, waited
assert_key_reported(locked, api_key, DAY, key_metadata(user=owner, email=email), seeded_metrics(1))
assert_key_reported(unlocked, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1))
def test_concurrent_reads_over_every_route_all_name_a_fresh_key(gateway: Gateway) -> None:
api_key: Final = digest_no_key_table_holds()
alias: Final = _alias()
entity: Final = f"integration-entity-{uuid.uuid4().hex}"
entity_columns: Final = {route.table: route.entity_column for route in ROUTES if route.table != USER_SPEND}
with gateway.scenario() as scenario:
owner, email = user_with_an_email(scenario)
rows: Final = (
user_row(owner, api_key, DAY),
*(seeded_row(table, column, entity, api_key, DAY) for table, column in entity_columns.items()),
)
with (
daily_rows(rows),
spend_logs_of_key(api_key, (named_row(0, alias),)),
ThreadPoolExecutor(CONCURRENT_READS) as pool,
):
reads: Final = tuple(
pool.submit(_activity_on_route, gateway, ROUTES[index % len(ROUTES)], api_key, entity)
for index in range(CONCURRENT_READS)
)
responses: Final = tuple(read.result() for read in reads)
for response in responses:
assert_key_reported(
response, api_key, DAY, key_metadata(alias=alias, user=owner, email=email), seeded_metrics(1)
)

48
tests/integration/spend/test_daily_activity_key_owner_traffic.py
View file

@ -12,7 +12,7 @@ from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, Scenario, eventually
from integration._support.client import Gateway, Scenario, eventually, string_value
from integration._support.daily_activity import (
AGGREGATED_USER_ACTIVITY,
DAY,
@ -25,6 +25,7 @@ from integration._support.daily_activity import (
daily_rows,
key_metadata,
key_no_key_table_holds,
purge_key_from_the_key_tables,
seeded_metrics,
seeded_row,
user_row,
@ -42,6 +43,10 @@ REQUESTS_OF_KEY: Final = (
'SELECT COALESCE(SUM(api_requests), 0)::int AS requests FROM "LiteLLM_DailyUserSpend" '
"WHERE api_key=%s AND user_id=%s"
)
NAMED_SPEND_LOGS_OF_KEY: Final = (
'SELECT COUNT(*)::int AS named FROM "LiteLLM_SpendLogs" '
"WHERE api_key=%s AND NULLIF(metadata->>'user_api_key_alias', '') IS NOT NULL"
)
UNIFIED_ENDPOINTS: Final = ("/v1/chat/completions", "/v1/messages", "/v1/responses")
REQUESTS_OF_A_BURST: Final = 21
READS_DURING_A_BURST: Final = 30
@ -210,6 +215,14 @@ def _wait_for_requests(api_key: str, user: str, requests: int) -> None:
)
def _wait_for_named_spend_logs(api_key: str, requests: int) -> None:
eventually(
lambda: read_rows(NAMED_SPEND_LOGS_OF_KEY, (api_key,)),
lambda rows: rows[0]["named"] == requests,
seconds=70,
)
def _cli_session_token(user: str, team: str) -> str:
cli_user: Final = LiteLLM_UserTable(user_id=user, user_role="internal_user", teams=[team], models=[])
return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=cli_user, team_id=team, team_alias="cli-team")
@ -244,6 +257,39 @@ def test_key_used_on_every_unified_endpoint_is_reported_with_its_own_alias_and_u
)
def test_key_purged_from_the_key_tables_is_reported_with_the_alias_its_spend_logs_name(gateway: Gateway) -> None:
prompts: Final = (_prompt(), _prompt(), _prompt())
with wire_server(_provider) as wire, gateway.scenario() as scenario:
model: Final = _priced_model(scenario, wire.url)
owner, email = user_with_an_email(scenario)
alias: Final = f"integration-alias-{uuid.uuid4().hex}"
generated: Final = gateway.post("/key/generate", {"user_id": owner, "key_alias": alias, "models": [model]})
key: Final = string_value(generated["key"])
stored: Final = sha256(key.encode()).hexdigest()
try:
answers: Final = tuple(
gateway.request("POST", endpoint, _request_body(endpoint, model, prompt), key=key)
for endpoint, prompt in zip(UNIFIED_ENDPOINTS, prompts, strict=True)
)
assert [answer.status_code for answer in answers] == [200, 200, 200], [answer.text for answer in answers]
received: Final = _sent_for_callers(wire.drain())
assert [request.target for request in received] == [
"/v1/chat/completions",
"/v1/responses",
"/v1/responses",
]
_wait_for_requests(stored, owner, 3)
_wait_for_named_spend_logs(stored, 3)
finally:
purge_key_from_the_key_tables(stored)
assert_key_owner_and_totals(
_activity_around_today(gateway, stored),
stored,
key_metadata(alias=alias, user=owner, email=email, exists=False),
_totals_of_requests(3),
)
def test_cli_session_spend_is_reported_with_the_user_and_team_of_the_session(
gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None: