test: restore router usage database regression coverage

2026-10-01 02:02:20 +00:00 · 2026-09-24 18:55:30 -07:00 · 2026-09-24 18:55:30 -07:00 · bdbbc94fe3
commit bdbbc94fe3
parent ef1c76e4ad
1 changed files with 101 additions and 0 deletions
--- a/tests/proxy_behavior/management/test_auto_router_usage.py
+++ b/tests/proxy_behavior/management/test_auto_router_usage.py
@ -0,0 +1,101 @@
+import json
+from typing import Final
+from uuid import uuid4
+
+import pytest
+from httpx import AsyncClient
+
+from litellm.proxy.utils import PrismaClient
+
+from .actors import Actor, World
+
+pytestmark: Final = pytest.mark.asyncio(loop_scope="session")
+
+
+async def test_routing_usage_groups_spend_and_keeps_other_users_and_internal_calls_out(
+    proxy_client: AsyncClient, prisma: PrismaClient, world: World
+) -> None:
+    db: Final = prisma.db
+    caller: Final = world.keys[Actor.INTERNAL_USER]
+    owner: Final = caller.user_id
+    key: Final = f"routing-usage-{uuid4()}"
+    records: Final = (
+        ("fast", 1.0, "a", "SIMPLE", None, owner, "2026-01-01T00:00:00"),
+        ("fast", 1.5, "a", "SIMPLE", None, owner, "2026-01-01T12:00:00"),
+        ("strong", 5.0, "a", "SIMPLE", None, owner, "2026-01-01T12:00:00"),
+        ("fast", 0.5, "a", None, None, owner, "2026-01-01T12:00:00"),
+        ("fast", 2.0, "b", "COMPLEX", None, owner, "2026-01-01T12:00:00"),
+        ("fast", 3.0, None, None, None, owner, "2026-01-01T12:00:00"),
+        ("fast", 90.0, "a", "SIMPLE", "shadow_eval_router", owner, "2026-01-01T12:00:00"),
+        ("fast", 80.0, None, None, "autorouter_classifier", owner, "2026-01-01T12:00:00"),
+        ("fast", 70.0, "a", "SIMPLE", None, "other-user", "2026-01-01T12:00:00"),
+        ("fast", 60.0, "a", "SIMPLE", None, owner, "2026-01-02T00:00:00"),
+    )
+    try:
+        for index, (model, spend, router, tier, origin, user, at) in enumerate(records):
+            metadata: Final = json.dumps(
+                {
+                    "routing_decision": {"router_model_name": router, "router_type": "complexity", "tier": tier}
+                    if router
+                    else None,
+                    "internal_call_origin": origin,
+                }
+            )
+            await db.execute_raw(
+                'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, model, model_group, spend, '
+                '"user", "startTime", "endTime", metadata) '
+                "VALUES ($1, 'completion', $2, $3, 'renamed-alias', $4, $5, $6::timestamp, $6::timestamp, $7::jsonb)",
+                f"{key}-{index}",
+                key,
+                model,
+                spend,
+                user,
+                at,
+                metadata,
+            )
+        query: Final = {"start_date": "2026-01-01", "end_date": "2026-01-01", "api_key": key}
+        headers: Final = {"Authorization": f"Bearer {caller.cleartext}"}
+        model_usage: Final = await proxy_client.get(
+            "/auto_router/usage", params={**query, "destination_model": "fast"}, headers=headers
+        )
+        assert model_usage.status_code == 200, model_usage.text
+        assert model_usage.json() == [
+            {
+                "model": "fast", "router_name": router, "router_type": router_type,
+                "tier": tier, "requests": requests, "spend": spend,
+            }
+            for router, router_type, tier, requests, spend in (
+                (None, None, None, 1, 3.0),
+                ("a", "complexity", "SIMPLE", 2, 2.5),
+                ("b", "complexity", "COMPLEX", 1, 2.0),
+                ("a", "complexity", None, 1, 0.5),
+            )
+        ]
+        router_usage: Final = await proxy_client.get(
+            "/auto_router/usage", params={**query, "router_name": "a", "router_type": "complexity"}, headers=headers
+        )
+        assert router_usage.status_code == 200, router_usage.text
+        assert router_usage.json() == [
+            {
+                "model": model, "router_name": "a", "router_type": "complexity",
+                "tier": tier, "requests": requests, "spend": spend,
+            }
+            for model, tier, requests, spend in (
+                ("strong", "SIMPLE", 1, 5.0), ("fast", "SIMPLE", 2, 2.5), ("fast", None, 1, 0.5)
+            )
+        ]
+        empty: Final = await proxy_client.get(
+            "/auto_router/usage",
+            params={**query, "destination_model": "fast", "api_key": f"{key}-other"},
+            headers=headers,
+        )
+        assert empty.status_code == 200, empty.text
+        assert empty.json() == []
+        forbidden: Final = await proxy_client.get(
+            "/auto_router/usage",
+            params={**query, "destination_model": "fast", "user_id": "other-user"},
+            headers=headers,
+        )
+        assert forbidden.status_code == 403, forbidden.text
+    finally:
+        await db.execute_raw('DELETE FROM "LiteLLM_SpendLogs" WHERE api_key = $1', key)