scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-29 01:42:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

test(integration): restore the shared retention setting, capture seeded days once and kill a listening worker

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-23 09:27:25 +00:00
parent 4290a4d651
commit 74fca4514f
1 changed files with 48 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

67
tests/integration/spend/test_daily_tag_spend_retention.py
View file

@ -1,3 +1,4 @@
import json
import os
import signal
import uuid
@ -13,7 +14,7 @@ from pydantic import JsonValue, TypeAdapter
from tests.integration._support.client import Gateway, eventually, string_value
from tests.integration._support.database import read_rows
from tests.integration._support.process import owned_proxy, owned_proxy_process
from tests.integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process
CLEANUP_EVERY_MINUTE: Final = "* * * * *"
RETENTION_SETTING: Final = "maximum_daily_tag_spend_retention_period"
@ -58,14 +59,38 @@ def _spend_log_present(request_id: str) -> bool:
return bool(read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)))
def _forget_stored_retention_setting() -> None:
def _stored_retention_setting() -> JsonValue:
rows: Final = read_rows(
'SELECT param_value -> %s AS value FROM "LiteLLM_Config" WHERE param_name = %s',
(RETENTION_SETTING, "general_settings"),
)
return rows[0]["value"] if rows else None
def _store_retention_setting(value: JsonValue) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
if value is None:
connection.execute(
'UPDATE "LiteLLM_Config" SET param_value = param_value - %s WHERE param_name = %s',
(RETENTION_SETTING, "general_settings"),
)
return
connection.execute(
'UPDATE "LiteLLM_Config" SET param_value = param_value - %s WHERE param_name = %s',
(RETENTION_SETTING, "general_settings"),
'UPDATE "LiteLLM_Config" SET param_value = jsonb_set(param_value, ARRAY[%s], %s::jsonb) '
"WHERE param_name = %s",
(RETENTION_SETTING, json.dumps(value), "general_settings"),
)
def _listening_workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]:
port: Final = owned.gateway.client.base_url.port
return tuple(
child
for child in psutil.Process(owned.process.pid).children(recursive=True)
if any(conn.status == psutil.CONN_LISTEN and conn.laddr.port == port for conn in child.net_connections("inet"))
)
def _listed_retention_value(gateway: Gateway) -> JsonValue:
listed: Final = _SETTINGS.validate_json(
gateway.request("GET", "/config/list", params={"config_type": "general_settings"}).content
@ -117,7 +142,8 @@ def test_config_update_turns_on_daily_tag_spend_cleanup_without_a_restart(gatewa
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
expired, yesterday_of_cutoff, on_the_cutoff, today = _day(200), _day(31), _day(30), _day(0)
_seed_daily_tag_spend(tag, (expired, yesterday_of_cutoff, on_the_cutoff, today))
_forget_stored_retention_setting()
previously_stored: Final = _stored_retention_setting()
_store_retention_setting(None)
try:
config: Final = _cleanup_config(tmp_path, {})
with owned_proxy(gateway, tmp_path, {}, config=config, workers=2) as owned, owned.scenario() as scenario:
@ -133,7 +159,7 @@ def test_config_update_turns_on_daily_tag_spend_cleanup_without_a_restart(gatewa
assert remaining == (on_the_cutoff, today), remaining
assert _completion_id(owned, model).startswith("chatcmpl-")
finally:
_forget_stored_retention_setting()
_store_retention_setting(previously_stored)
_delete_daily_tag_spend(tag)
@ -143,7 +169,8 @@ def test_unparseable_daily_tag_spend_retention_deletes_nothing_and_keeps_serving
) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
request_id: Final = f"integration-retention-{uuid.uuid4().hex}"
_seed_daily_tag_spend(tag, (_day(200),))
expired: Final = _day(200)
_seed_daily_tag_spend(tag, (expired,))
_seed_old_spend_log(request_id, days_ago=200)
try:
config: Final = _cleanup_config(
@ -152,7 +179,7 @@ def test_unparseable_daily_tag_spend_retention_deletes_nothing_and_keeps_serving
with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario:
model: Final = scenario.model()
eventually(lambda: _spend_log_present(request_id), lambda present: not present, seconds=150)
assert _remaining_days(tag) == (_day(200),)
assert _remaining_days(tag) == (expired,)
assert _completion_id(owned, model).startswith("chatcmpl-")
finally:
_delete_daily_tag_spend(tag)
@ -164,7 +191,8 @@ def test_daily_tag_spend_keeps_days_the_shorter_spend_log_horizon_already_pruned
) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
request_id: Final = f"integration-retention-{uuid.uuid4().hex}"
_seed_daily_tag_spend(tag, (_day(200), _day(60)))
expired, inside_tag_horizon = _day(200), _day(60)
_seed_daily_tag_spend(tag, (expired, inside_tag_horizon))
_seed_old_spend_log(request_id, days_ago=60)
try:
config: Final = _cleanup_config(
@ -172,8 +200,8 @@ def test_daily_tag_spend_keeps_days_the_shorter_spend_log_horizon_already_pruned
)
with owned_proxy(gateway, tmp_path, {}, config=config):
eventually(lambda: _spend_log_present(request_id), lambda present: not present, seconds=150)
remaining: Final = eventually(lambda: _remaining_days(tag), lambda days: _day(200) not in days, seconds=150)
assert remaining == (_day(60),), remaining
remaining: Final = eventually(lambda: _remaining_days(tag), lambda days: expired not in days, seconds=150)
assert remaining == (inside_tag_horizon,), remaining
finally:
_delete_daily_tag_spend(tag)
@ -181,24 +209,24 @@ def test_daily_tag_spend_keeps_days_the_shorter_spend_log_horizon_already_pruned
@pytest.mark.covers("spend.daily_tag_spend.cleanup_and_serving_survive_losing_one_of_two_workers")
def test_daily_tag_spend_cleanup_completes_after_one_of_two_workers_is_killed(gateway: Gateway, tmp_path: Path) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
_seed_daily_tag_spend(tag, (_day(200), _day(0)))
expired, today = _day(200), _day(0)
_seed_daily_tag_spend(tag, (expired, today))
try:
config: Final = _cleanup_config(tmp_path, {RETENTION_SETTING: "30d"})
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
with owned.gateway.scenario() as scenario:
model: Final = scenario.model()
workers: Final = eventually(
lambda: psutil.Process(owned.process.pid).children(recursive=True),
lambda children: len(children) >= 2,
seconds=30,
lambda: _listening_workers(owned), lambda found: len(found) == 2, seconds=30
)
workers[0].send_signal(signal.SIGKILL)
eventually(lambda: workers[0].is_running(), lambda alive: not alive, seconds=10)
ids: Final = tuple(_completion_id(owned.gateway, model) for _ in range(6))
assert len(set(ids)) == 6 and all(identity.startswith("chatcmpl-") for identity in ids), ids
remaining: Final = eventually(
lambda: _remaining_days(tag), lambda days: _day(200) not in days, seconds=150
lambda: _remaining_days(tag), lambda days: expired not in days, seconds=150
)
assert remaining == (_day(0),), remaining
assert remaining == (today,), remaining
finally:
_delete_daily_tag_spend(tag)
@ -207,12 +235,13 @@ def test_daily_tag_spend_cleanup_completes_after_one_of_two_workers_is_killed(ga
def test_daily_tag_spend_is_kept_forever_when_its_retention_is_unset(gateway: Gateway, tmp_path: Path) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
request_id: Final = f"integration-retention-{uuid.uuid4().hex}"
_seed_daily_tag_spend(tag, (_day(200),))
expired: Final = _day(200)
_seed_daily_tag_spend(tag, (expired,))
_seed_old_spend_log(request_id, days_ago=200)
try:
config: Final = _cleanup_config(tmp_path, {"maximum_spend_logs_retention_period": "30d"})
with owned_proxy(gateway, tmp_path, {}, config=config):
eventually(lambda: _spend_log_present(request_id), lambda present: not present, seconds=150)
assert _remaining_days(tag) == (_day(200),)
assert _remaining_days(tag) == (expired,)
finally:
_delete_daily_tag_spend(tag)