mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
test(alerts): verify resumable work with complete isolated fixtures
This commit is contained in:
moe-berri
parent
063df89cb9
commit
1f535c74a7
2 changed files with 46 additions and 32 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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 + "%")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue