From 1f535c74a787d6ce356b451ca666be23f181623e Mon Sep 17 00:00:00 2001
From: moe-berri <moe@berri.ai>
Date: Sun, 8 Nov 2026 18:30:36 +0800
Subject: [PATCH] test(alerts): verify resumable work with complete isolated
 fixtures

---
 .../database/test_performance_issues.py       | 36 +++++++++++-----
 .../database/test_provider_latency.py         | 42 +++++++++----------
 2 files changed, 46 insertions(+), 32 deletions(-)

diff --git a/tests/integration/database/test_performance_issues.py b/tests/integration/database/test_performance_issues.py
index 6439212aad7..930ae4dd167 100644
--- a/tests/integration/database/test_performance_issues.py
+++ b/tests/integration/database/test_performance_issues.py
@@ -22,7 +22,13 @@ from litellm.repositories.daily_activity_buckets import api_key_bucket
 from litellm.repositories.issues import IssuesRepository
 from litellm.repositories.usage_alert_metrics import IssueMetricWindow, UsageAlertMetricsRepository
 from litellm.repositories.usage_alerts import UsageAlertsRepository
-from litellm.types.proxy.usage_alerts import UsageAlertObservation, AlertRule, AlertRuleInput, UsageAlertScope
+from litellm.types.proxy.usage_alerts import (
+    AlertRule,
+    AlertRuleInput,
+    UsageAlertObservation,
+    UsageAlertProjectionWork,
+    UsageAlertScope,
+)
 
 _DAY: Final = date(2026, 11, 4)
 _NOW: Final = datetime(2026, 11, 5, 3, microsecond=987654, tzinfo=timezone.utc)
@@ -185,6 +191,14 @@ async def _coverage(prisma: PrismaClient, day: date) -> None:
     )
 
 
+async def _finish_projection(job: UsageAlertJob, now: datetime, work: UsageAlertProjectionWork) -> None:
+    async with asyncio.timeout(60):
+        while not await job._performance_projection.run_pending(
+            now, work, job._store, asyncio.get_running_loop().time() + 10
+        ):
+            continue
+
+
 async def _seed(prisma: PrismaClient, *, reliability: bool = False) -> None:
     for table, entity in (("LiteLLM_DailyUserSpend", "user_id"), ("LiteLLM_DailyTeamSpend", "team_id")):
         await prisma.writer_db.query_raw(
@@ -375,8 +389,14 @@ async def test_retirement_detects_used_model_after_unused_model_in_the_same_alia
 async def test_single_key_reliability_collapses_matching_scopes_without_traffic_growth(db: ManagementDatabase) -> None:
     async with _isolated(db) as prisma:
         await _seed(prisma, reliability=True)
-        await UsageAlertJob.from_prisma(prisma, SlackAlerting(alerting=[], alert_types=[]))._run(_NOW)
+        job: Final = UsageAlertJob.from_prisma(prisma, SlackAlerting(alerting=[], alert_types=[]))
         issues: Final = IssuesRepository.from_prisma(prisma)
+        async with asyncio.timeout(60):
+            while True:
+                await job.run(_NOW)
+                if await job._store.pending_projection_work() is None and (await issues.list_issues()).total:
+                    break
+                await asyncio.sleep(0.1)
         page: Final = await issues.list_issues()
         assert page.total == 1, tuple((item.cohort.scope_type, item.cohort.scope_id, item.title) for item in page.data)
         issue: Final = page.data[0]
@@ -632,19 +652,15 @@ async def test_entity_page_progress_resumes_old_day_without_early_completion(db:
         assert "evaluation_progress" not in AlertRule.model_json_schema(mode="serialization")["properties"]
         first: Final = await job._store.list_alerts(limit=500)
         assert first.total == 10
-        projector: Final = job._performance_projection
-        assert projector is not None
         work: Final = await job._store.pending_projection_work()
         assert work is not None
-        assert await projector.run_pending(_NOW, work, job._store, asyncio.get_running_loop().time() + 10)
+        await _finish_projection(job, _NOW, work)
         await job._store.finish_projection_work(work)
         assert (await IssuesRepository.from_prisma(prisma).list_issues()).total == 10
         await job._scope_page((partial,), _DAY, True)
         next_work: Final = await job._store.pending_projection_work()
         assert next_work is not None and next_work.version != work.version
-        assert await projector.run_pending(
-            _NOW + timedelta(hours=1), next_work, job._store, asyncio.get_running_loop().time() + 10
-        )
+        await _finish_projection(job, _NOW + timedelta(hours=1), next_work)
         await job._store.finish_projection_work(next_work)
         assert (await IssuesRepository.from_prisma(prisma).list_issues()).total == 20
         await job._run(_NOW + timedelta(days=1))
@@ -842,11 +858,9 @@ async def test_partial_broad_default_waits_for_its_key_family_only(db: Managemen
         await job._scope_page(spend, _DAY, True)
         assert await sources.active_signals(_DAY, scope_type="key", metric="spend", completed_only=True)
         assert not await sources.active_signals(_DAY, scope_type="deployment", metric="spend", completed_only=True)
-        projector: Final = job._performance_projection
-        assert projector is not None
         work: Final = await job._store.pending_projection_work()
         assert work is not None
-        assert await projector.run_pending(_NOW, work, job._store, asyncio.get_running_loop().time() + 10)
+        await _finish_projection(job, _NOW, work)
         partial: Final = await sources.list_issues()
         assert len(partial.data) == 1 and partial.data[0].cohort.scope_type == "key"
         while not await job._store.completed_rule_ids((spend[0],), _DAY):
diff --git a/tests/integration/database/test_provider_latency.py b/tests/integration/database/test_provider_latency.py
index 387f40858da..b63075b7a4b 100644
--- a/tests/integration/database/test_provider_latency.py
+++ b/tests/integration/database/test_provider_latency.py
@@ -3,6 +3,7 @@ from uuid import uuid4
 
 import pytest
 from integration._support.actors import ManagementDatabase
+from integration.database.test_performance_issues import _coverage, _isolated
 
 from litellm.repositories.provider_latency import ProviderLatencyRepository
 from litellm.types.proxy.issues import AlertModelPair
@@ -61,9 +62,9 @@ async def test_provider_latency_requires_established_within_model_regressions(db
     pairs = tuple(AlertModelPair(provider=prefix, model=f"{prefix}-{suffix}") for suffix in ("a", "b"))
     keys = (prefix + "-key-a", prefix + "-key-b")
     rule = AlertRuleInput(name="Latency", scope_type="key", metric="latency", direction="increase", min_value=500)
-    repository = ProviderLatencyRepository(db.prisma.writer_db)
-    try:
-        await db.prisma.writer_db.query_raw(
+    async with _isolated(db) as prisma:
+        repository = ProviderLatencyRepository(prisma.writer_db)
+        await prisma.writer_db.query_raw(
             """INSERT INTO "LiteLLM_DailyUserSpend"
                 (id,user_id,date,api_key,model,custom_llm_provider,api_requests,timed_requests,total_response_time_ms,updated_at)
             SELECT $1||d||m.suffix||k.suffix,$1||'-owner',($2::date-d)::text,$1||'-key-'||k.suffix,
@@ -73,7 +74,7 @@ async def test_provider_latency_requires_established_within_model_regressions(db
             prefix,
             day.isoformat(),
         )
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             """INSERT INTO "LiteLLM_DailyTeamSpend"
                 (id,team_id,date,api_key,model,custom_llm_provider,api_requests,timed_requests,total_response_time_ms,updated_at)
             SELECT id||'team',$1||CASE WHEN api_key=$1||'-key-a' THEN 'team1' ELSE 'team2' END,
@@ -104,7 +105,7 @@ async def test_provider_latency_requires_established_within_model_regressions(db
         teams = (prefix + "team1", prefix + "team2")
         assert await repository.source_members(pairs, day, rule, "team", teams) == frozenset(teams)
         assert await repository.source_members(pairs, day, rule, "deployment", ("deployment",)) == {"deployment"}
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             """INSERT INTO "LiteLLM_DailyUserSpend"
                 (id,user_id,date,api_key,model,custom_llm_provider,api_requests,timed_requests,total_response_time_ms,updated_at)
             SELECT id||'-new-owner',$1||'-new-owner',date,api_key,model,custom_llm_provider,api_requests,
@@ -116,24 +117,26 @@ async def test_provider_latency_requires_established_within_model_regressions(db
         assert await repository.source_members(
             pairs, day, rule, "user", (prefix + "-owner", prefix + "-new-owner")
         ) == {prefix + "-owner"}
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'DELETE FROM "LiteLLM_DailyUserSpend" WHERE user_id=$1 RETURNING id', prefix + "-new-owner"
         )
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms=total_response_time_ms/5 '
             "WHERE api_key=$1 AND date=$2 RETURNING id",
             keys[1],
             day.isoformat(),
         )
         assert await repository.latency_cohort(pairs, day, rule) is None
+        assert not await repository.cohort_recovered(pairs, day, rule)
+        await _coverage(prisma, day)
         assert await repository.cohort_recovered(pairs, day, rule)
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms=total_response_time_ms*5 '
             "WHERE api_key=$1 AND date=$2 RETURNING id",
             keys[1],
             day.isoformat(),
         )
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             """INSERT INTO "LiteLLM_DailyUserSpend"
                 (id,user_id,date,api_key,model,custom_llm_provider,api_requests,timed_requests,total_response_time_ms,updated_at)
             SELECT $1||'new'||m,$1||'-new-owner',$2,$1||'-new-key',$1||'-'||m,$1,50,50,2500000,CURRENT_TIMESTAMP
@@ -147,15 +150,15 @@ async def test_provider_latency_requires_established_within_model_regressions(db
         assert (with_new_key.total_timed_requests, with_new_key.comparable_timed_requests) == (500, 400)
         assert (with_new_key.impact.keys, with_new_key.impact.users, with_new_key.impact.requests) == (3, 2, 500)
         assert with_new_key.impact.measured_keys == 2
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'UPDATE "LiteLLM_DailyUserSpend" SET timed_requests=51 WHERE api_key=$1 RETURNING id', prefix + "-new-key"
         )
         assert await repository.latency_cohort(pairs, day, rule) is None
         assert not await repository.cohort_recovered(pairs, day, rule)
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'DELETE FROM "LiteLLM_DailyUserSpend" WHERE api_key=$1 RETURNING id', prefix + "-new-key"
         )
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             """INSERT INTO "LiteLLM_DailyUserSpend"
                 (id,user_id,date,api_key,model,custom_llm_provider,api_requests,timed_requests,total_response_time_ms,updated_at)
             SELECT id||'-other',user_id,date,api_key,model,$1||'-other',1000,1000,100000,CURRENT_TIMESTAMP
@@ -169,7 +172,7 @@ async def test_provider_latency_requires_established_within_model_regressions(db
         assert await repository.source_members(pairs, day, rule, "user", (prefix + "-owner",)) == frozenset()
         assert await repository.source_members(pairs, day, rule, "deployment", ("deployment",)) == frozenset()
         assert await repository.source_members(pairs, day, rule, "team", teams) == frozenset(teams)
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             """INSERT INTO "LiteLLM_DailyTeamSpend"
                 (id,team_id,date,api_key,model,custom_llm_provider,api_requests,timed_requests,total_response_time_ms,updated_at)
             SELECT id||'-other',team_id,date,api_key,model,$1||'-other',1000,1000,100000,CURRENT_TIMESTAMP
@@ -178,7 +181,7 @@ async def test_provider_latency_requires_established_within_model_regressions(db
             day.isoformat(),
         )
         assert await repository.source_members(pairs, day, rule, "team", teams) == frozenset()
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             """UPDATE "LiteLLM_DailyUserSpend"
                SET api_requests=CASE WHEN model=$1||'-a' AND date<$2 THEN 10000
                                      WHEN model=$1||'-b' AND date=$2 THEN 5000 ELSE 100 END,
@@ -193,34 +196,31 @@ async def test_provider_latency_requires_established_within_model_regressions(db
         )
         assert await repository.latency_cohort(pairs, day, rule) is None
         assert await repository.cohort_recovered(pairs, day, rule)
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms=total_response_time_ms*5 '
             "WHERE custom_llm_provider=$1 AND date=$2 RETURNING id",
             prefix,
             day.isoformat(),
         )
         assert await repository.latency_cohort(pairs, day, rule) is not None
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms=0 WHERE custom_llm_provider=$1 AND date<$2 RETURNING id',
             prefix,
             day.isoformat(),
         )
         assert await repository.latency_cohort(pairs, day, rule) is None
         assert not await repository.cohort_recovered(pairs, day, rule)
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms=timed_requests*'
             "(CASE WHEN model=$1||'-a' THEN 200 ELSE 1000 END) "
             "WHERE custom_llm_provider=$1 AND date<$2 RETURNING id",
             prefix,
             day.isoformat(),
         )
-        await db.prisma.writer_db.query_raw(
+        await prisma.writer_db.query_raw(
             'DELETE FROM "LiteLLM_DailyUserSpend" WHERE custom_llm_provider=$1 AND date<($2::date-2)::text RETURNING id',
             prefix,
             day.isoformat(),
         )
         assert await repository.latency_cohort(pairs, day, rule) is None
         assert not await repository.cohort_recovered(pairs, day, rule)
-    finally:
-        for table in ('"LiteLLM_DailyUserSpend"', '"LiteLLM_DailyTeamSpend"'):
-            await db.prisma.writer_db.query_raw(f"DELETE FROM {table} WHERE id LIKE $1 RETURNING id", prefix + "%")