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>

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)