From 486465808169426ade8d867ec53dbd6b52db31f0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 09:27:25 +0000 Subject: [PATCH] test(integration): restore the shared retention setting, capture seeded days once and kill a listening worker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend/test_daily_tag_spend_retention.py | 67 +++++++++++++------ 1 file changed, 48 insertions(+), 19 deletions(-) diff --git a/tests/integration/spend/test_daily_tag_spend_retention.py b/tests/integration/spend/test_daily_tag_spend_retention.py index 475d7ecde11..8642d073941 100644 --- a/tests/integration/spend/test_daily_tag_spend_retention.py +++ b/tests/integration/spend/test_daily_tag_spend_retention.py @@ -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)