From 04f63dc235db66cb276291063382a077e577fba7 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 08:54:02 +0000 Subject: [PATCH] test(integration): cover runtime update, invalid value, independent horizons and worker loss for daily tag spend retention Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend/test_daily_tag_spend_retention.py | 119 +++++++++++++++++- 1 file changed, 117 insertions(+), 2 deletions(-) diff --git a/tests/integration/spend/test_daily_tag_spend_retention.py b/tests/integration/spend/test_daily_tag_spend_retention.py index 5cbee96a2b1..475d7ecde11 100644 --- a/tests/integration/spend/test_daily_tag_spend_retention.py +++ b/tests/integration/spend/test_daily_tag_spend_retention.py @@ -1,20 +1,24 @@ import os +import signal import uuid from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Final +import psutil import psycopg import pytest import yaml from pydantic import JsonValue, TypeAdapter -from tests.integration._support.client import Gateway, eventually +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 +from tests.integration._support.process import owned_proxy, owned_proxy_process CLEANUP_EVERY_MINUTE: Final = "* * * * *" +RETENTION_SETTING: Final = "maximum_daily_tag_spend_retention_period" _MAPPING: Final = TypeAdapter(dict[str, JsonValue]) +_SETTINGS: Final = TypeAdapter(list[dict[str, JsonValue]]) def _day(days_ago: int) -> str: @@ -54,6 +58,26 @@ 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: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute( + 'UPDATE "LiteLLM_Config" SET param_value = param_value - %s WHERE param_name = %s', + (RETENTION_SETTING, "general_settings"), + ) + + +def _listed_retention_value(gateway: Gateway) -> JsonValue: + listed: Final = _SETTINGS.validate_json( + gateway.request("GET", "/config/list", params={"config_type": "general_settings"}).content + ) + matching: Final = tuple(entry for entry in listed if entry["field_name"] == RETENTION_SETTING) + return matching[0]["field_value"] if matching else "not listed" + + +def _completion_id(gateway: Gateway, model: str) -> str: + return string_value(gateway.chat(model, text=f"retention audit {uuid.uuid4().hex}")["id"]) + + def _cleanup_config(tmp_path: Path, retention: dict[str, JsonValue]) -> Path: base: Final = _MAPPING.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = { @@ -88,6 +112,97 @@ def test_daily_tag_spend_retention_prunes_only_rows_older_than_the_period(gatewa _delete_daily_tag_spend(tag) +@pytest.mark.covers("spend.daily_tag_spend.runtime_config_update_enables_cleanup_on_a_multi_worker_proxy") +def test_config_update_turns_on_daily_tag_spend_cleanup_without_a_restart(gateway: Gateway, tmp_path: Path) -> None: + 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() + try: + config: Final = _cleanup_config(tmp_path, {}) + with owned_proxy(gateway, tmp_path, {}, config=config, workers=2) as owned, owned.scenario() as scenario: + model: Final = scenario.model() + assert _listed_retention_value(owned) is None + owned.post("/config/update", {"general_settings": {RETENTION_SETTING: "30d"}}) + assert _listed_retention_value(owned) == "30d" + remaining: Final = eventually( + lambda: _remaining_days(tag), + lambda days: yesterday_of_cutoff not in days, + seconds=150, + ) + assert remaining == (on_the_cutoff, today), remaining + assert _completion_id(owned, model).startswith("chatcmpl-") + finally: + _forget_stored_retention_setting() + _delete_daily_tag_spend(tag) + + +@pytest.mark.covers("spend.daily_tag_spend.invalid_retention_value_keeps_rows_and_leaves_the_proxy_serving") +def test_unparseable_daily_tag_spend_retention_deletes_nothing_and_keeps_serving( + 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),)) + _seed_old_spend_log(request_id, days_ago=200) + try: + config: Final = _cleanup_config( + tmp_path, {RETENTION_SETTING: "soon", "maximum_spend_logs_retention_period": "30d"} + ) + 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 _completion_id(owned, model).startswith("chatcmpl-") + finally: + _delete_daily_tag_spend(tag) + + +@pytest.mark.covers("spend.daily_tag_spend.retention_horizon_is_independent_of_the_spend_log_horizon") +def test_daily_tag_spend_keeps_days_the_shorter_spend_log_horizon_already_pruned( + 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), _day(60))) + _seed_old_spend_log(request_id, days_ago=60) + try: + config: Final = _cleanup_config( + tmp_path, {RETENTION_SETTING: "90d", "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) + remaining: Final = eventually(lambda: _remaining_days(tag), lambda days: _day(200) not in days, seconds=150) + assert remaining == (_day(60),), remaining + finally: + _delete_daily_tag_spend(tag) + + +@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))) + 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, + ) + workers[0].send_signal(signal.SIGKILL) + 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 + ) + assert remaining == (_day(0),), remaining + finally: + _delete_daily_tag_spend(tag) + + @pytest.mark.covers("spend.daily_tag_spend.unset_retention_never_deletes_even_while_spend_logs_are_pruned") 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}"