From ef9385c712297fa59bf8555962bfd719c55a4529 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 30 Oct 2026 16:25:11 -0700 Subject: [PATCH] refactor(proxy): inject gateway activity through a typed repository (#49116) * refactor(proxy): inject gateway activity through a typed repository * ci: keep dependency tests in their existing shard --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/models/gateway_activity.py | 31 +++++ litellm/proxy/deps/repositories.py | 11 ++ .../gateway_request_endpoints.py | 74 ++---------- .../gateway_activity_repository.py | 41 +++++++ tests/unit/proxy/deps/test_repositories.py | 29 ++++- .../test_gateway_request_endpoints.py | 108 ++++++++---------- .../test_gateway_activity_repository.py | 47 ++++++++ 7 files changed, 215 insertions(+), 126 deletions(-) create mode 100644 litellm/models/gateway_activity.py create mode 100644 litellm/repositories/gateway_activity_repository.py create mode 100644 tests/unit/repositories/test_gateway_activity_repository.py diff --git a/litellm/models/gateway_activity.py b/litellm/models/gateway_activity.py new file mode 100644 index 00000000000..a37e4e66817 --- /dev/null +++ b/litellm/models/gateway_activity.py @@ -0,0 +1,31 @@ +from typing import Literal, Protocol, TypeAlias + +from pydantic import ConfigDict + +from litellm.types.llms.base import LiteLLMBaseModel + + +class GatewayRouteActivity(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["route"] + date: str + category: str + route: str + successful_requests: int + failed_requests: int + + +class GatewayStatusCodeActivity(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["status"] + status_code: int + failed_requests: int + + +GatewayActivity: TypeAlias = GatewayRouteActivity | GatewayStatusCodeActivity + + +class GatewayActivityStore(Protocol): + async def get_activity(self, start_date: str, end_date: str) -> tuple[GatewayActivity, ...]: ... diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index f5308f837ae..e37ecfd582a 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -3,9 +3,11 @@ from typing import Annotated, TypeAlias from fastapi import Depends from litellm.models.deployment_warning import DeploymentWarningIgnoreStore +from litellm.models.gateway_activity import GatewayActivityStore from litellm.models.user_banner import UserBannerStore from litellm.proxy.deps.database import DatabaseDep from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository +from litellm.repositories.gateway_activity_repository import GatewayActivityRepository from litellm.repositories.user_banner_repository import UserBannerRepository @@ -23,3 +25,12 @@ async def get_deployment_warning_ignore_repository(database: DatabaseDep) -> Dep DeploymentWarningIgnoreRepositoryDep: TypeAlias = Annotated[ DeploymentWarningIgnoreStore | None, Depends(get_deployment_warning_ignore_repository) ] + + +async def get_gateway_activity_repository(database: DatabaseDep) -> GatewayActivityStore | None: + return None if database is None else GatewayActivityRepository(database.prisma) + + +GatewayActivityRepositoryDep: TypeAlias = Annotated[ + GatewayActivityStore | None, Depends(get_gateway_activity_repository) +] diff --git a/litellm/proxy/management_endpoints/gateway_request_endpoints.py b/litellm/proxy/management_endpoints/gateway_request_endpoints.py index 4052e6f714e..01c9ad137c2 100644 --- a/litellm/proxy/management_endpoints/gateway_request_endpoints.py +++ b/litellm/proxy/management_endpoints/gateway_request_endpoints.py @@ -13,15 +13,15 @@ and the endpoint is restricted to proxy admin roles. from collections.abc import Sequence from datetime import datetime, timedelta, timezone -from typing import Annotated, Final, Literal +from typing import Annotated, Final from fastapi import APIRouter, Depends, HTTPException, Query -from pydantic import ConfigDict, Field, TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.types.llms.base import LiteLLMBaseModel +from litellm.proxy.deps.repositories import GatewayActivityRepositoryDep from litellm.types.proxy.gateway_requests import ( GatewayRequestActivityResponse, GatewayRequestBreakdownEntry, @@ -33,50 +33,6 @@ router: Final = APIRouter() _DEFAULT_LOOKBACK_DAYS: Final = 30 -_ACTIVITY_SQL: Final = """ - SELECT - 'route' AS kind, date, category, route, - SUM(successful_requests)::bigint AS successful_requests, - SUM(failed_requests)::bigint AS failed_requests, - NULL::integer AS status_code - FROM "LiteLLM_DailyGatewayRequests" - WHERE date >= $1 AND date <= $2 - GROUP BY date, category, route - UNION ALL - SELECT - 'status' AS kind, NULL, NULL, NULL, - 0::bigint, - SUM(failed_requests)::bigint, - status_code - FROM "LiteLLM_DailyGatewayFailedRequests" - WHERE date >= $1 AND date <= $2 - GROUP BY status_code -""" - - -class _AggregateRow(LiteLLMBaseModel): - """Validates one query_raw row so the handler works with typed values, not Any.""" - - kind: Literal["route"] - date: str - category: str - route: str - successful_requests: int - failed_requests: int - - -class _StatusCodeAggregateRow(LiteLLMBaseModel): - model_config = ConfigDict(frozen=True) - - kind: Literal["status"] - status_code: int - failed_requests: int - - -_ACTIVITY_ROWS_ADAPTER: Final = TypeAdapter( - tuple[Annotated[_AggregateRow | _StatusCodeAggregateRow, Field(discriminator="kind")], ...] -) - def _default_range() -> tuple[str, str]: end: Final = datetime.now(timezone.utc) @@ -84,7 +40,7 @@ def _default_range() -> tuple[str, str]: return start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d") -def _fold_by_date(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestDailyEntry, ...]: +def _fold_by_date(rows: Sequence[GatewayRouteActivity]) -> tuple[GatewayRequestDailyEntry, ...]: dates: Final = sorted(frozenset(row.date for row in rows)) return tuple( GatewayRequestDailyEntry( @@ -96,7 +52,7 @@ def _fold_by_date(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestDailyEnt ) -def _fold_by_route(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestBreakdownEntry, ...]: +def _fold_by_route(rows: Sequence[GatewayRouteActivity]) -> tuple[GatewayRequestBreakdownEntry, ...]: pairs: Final = sorted(frozenset((row.category, row.route) for row in rows)) entries: Final = tuple( GatewayRequestBreakdownEntry( @@ -113,7 +69,7 @@ def _fold_by_route(rows: Sequence[_AggregateRow]) -> tuple[GatewayRequestBreakdo def _fold_by_status_code( - rows: Sequence[_StatusCodeAggregateRow], + rows: Sequence[GatewayStatusCodeActivity], ) -> tuple[GatewayRequestStatusCodeEntry, ...]: return tuple( GatewayRequestStatusCodeEntry(status_code=row.status_code, failed_requests=row.failed_requests) @@ -128,6 +84,7 @@ def _fold_by_status_code( ) async def get_gateway_daily_activity( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + activity: GatewayActivityRepositoryDep, start_date: str | None = Query(default=None, description="Start date in YYYY-MM-DD format"), end_date: str | None = Query(default=None, description="End date in YYYY-MM-DD format"), ) -> GatewayRequestActivityResponse: @@ -137,8 +94,6 @@ async def get_gateway_daily_activity( Deployment-wide: the underlying table has no per-key or per-user dimension, so this is admin-only. """ - from litellm.proxy.proxy_server import prisma_client - if user_api_key_dict.user_role not in ( LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, @@ -148,22 +103,15 @@ async def get_gateway_daily_activity( detail="Only proxy admin roles can view gateway request counts across the deployment", ) - if prisma_client is None: + if activity is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) default_start, default_end = _default_range() selected_start: Final = start_date or default_start selected_end: Final = end_date or default_end - raw_rows: Final = await prisma_client.db.query_raw( # pyright: ignore[reportAny] # untyped prisma client - _ACTIVITY_SQL, - selected_start, - selected_end, - ) - activity_rows: Final = _ACTIVITY_ROWS_ADAPTER.validate_python(raw_rows or ()) - rows: Final[tuple[_AggregateRow, ...]] = tuple(row for row in activity_rows if isinstance(row, _AggregateRow)) - status_code_rows: Final[tuple[_StatusCodeAggregateRow, ...]] = tuple( - row for row in activity_rows if isinstance(row, _StatusCodeAggregateRow) - ) + activity_rows: Final = await activity.get_activity(selected_start, selected_end) + rows: Final = tuple(row for row in activity_rows if isinstance(row, GatewayRouteActivity)) + status_code_rows: Final = tuple(row for row in activity_rows if isinstance(row, GatewayStatusCodeActivity)) verbose_proxy_logger.debug("/gateway/daily/activity - aggregated %d rows", len(rows)) return GatewayRequestActivityResponse( diff --git a/litellm/repositories/gateway_activity_repository.py b/litellm/repositories/gateway_activity_repository.py new file mode 100644 index 00000000000..2950668a931 --- /dev/null +++ b/litellm/repositories/gateway_activity_repository.py @@ -0,0 +1,41 @@ +from typing import TYPE_CHECKING, Annotated, Final + +from pydantic import Field, TypeAdapter + +from litellm.models.gateway_activity import GatewayActivity + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +_ACTIVITY_SQL: Final = """ + SELECT + 'route' AS kind, date, category, route, + SUM(successful_requests)::bigint AS successful_requests, + SUM(failed_requests)::bigint AS failed_requests, + NULL::integer AS status_code + FROM "LiteLLM_DailyGatewayRequests" + WHERE date >= $1 AND date <= $2 + GROUP BY date, category, route + UNION ALL + SELECT + 'status' AS kind, NULL, NULL, NULL, + 0::bigint, + SUM(failed_requests)::bigint, + status_code + FROM "LiteLLM_DailyGatewayFailedRequests" + WHERE date >= $1 AND date <= $2 + GROUP BY status_code +""" + +_ACTIVITY_ROWS_ADAPTER: Final = TypeAdapter(tuple[Annotated[GatewayActivity, Field(discriminator="kind")], ...]) + + +class GatewayActivityRepository: + def __init__(self, prisma_client: "PrismaClient") -> None: + self._prisma_client = prisma_client + + async def get_activity(self, start_date: str, end_date: str) -> tuple[GatewayActivity, ...]: + raw_rows: Final = await self._prisma_client.db.query_raw( # pyright: ignore[reportAny] # untyped Prisma client + _ACTIVITY_SQL, start_date, end_date + ) + return _ACTIVITY_ROWS_ADAPTER.validate_python(raw_rows or ()) diff --git a/tests/unit/proxy/deps/test_repositories.py b/tests/unit/proxy/deps/test_repositories.py index aa4af5310ff..e802a4fda61 100644 --- a/tests/unit/proxy/deps/test_repositories.py +++ b/tests/unit/proxy/deps/test_repositories.py @@ -8,9 +8,14 @@ from fastapi import FastAPI from fastapi.testclient import TestClient from litellm.models.deployment_warning import AdminIgnore +from litellm.models.gateway_activity import GatewayActivity from litellm.models.user_banner import UserBanner from litellm.proxy.deps.database import get_database -from litellm.proxy.deps.repositories import DeploymentWarningIgnoreRepositoryDep, UserBannerRepositoryDep +from litellm.proxy.deps.repositories import ( + DeploymentWarningIgnoreRepositoryDep, + GatewayActivityRepositoryDep, + UserBannerRepositoryDep, +) from litellm.repositories import database as database_runtime from tests._support.native_database import native_database @@ -32,6 +37,10 @@ def _app(database: database_runtime.Database | None) -> FastAPI: async def ignores(ignores: DeploymentWarningIgnoreRepositoryDep) -> Mapping[str, AdminIgnore] | None: return None if ignores is None else await ignores.get_ignores() + @app.get("/activity") + async def activity(activity: GatewayActivityRepositoryDep) -> tuple[GatewayActivity, ...] | None: + return None if activity is None else await activity.get_activity("2026-01-01", "2026-01-31") + return app @@ -83,3 +92,21 @@ def test_dependency_override_changes_next_request_without_changing_lifespan_owne assert client.get("/banner").json() == UserBanner(message="stored").model_dump() monkeypatch.setitem(app.dependency_overrides, get_database, lambda: None) assert client.get("/banner").json() == UserBanner().model_dump() + + +def test_gateway_activity_reads_the_lifespan_reader_not_the_writer() -> None: + prisma: Final = MagicMock() + row: Final = {"kind": "status", "status_code": 429, "failed_requests": 3} + prisma.db.query_raw = AsyncMock(return_value=[row]) + prisma.writer_db.query_raw = AsyncMock(return_value=[{**row, "failed_requests": 99}]) + with TestClient(_app(database_runtime.Database(prisma=prisma, native=NATIVE))) as client: + response: Final = client.get("/activity") + assert response.status_code == 200 + assert response.json() == [row] + + +def test_gateway_activity_without_database_is_unavailable() -> None: + with TestClient(_app(None)) as client: + response: Final = client.get("/activity") + assert response.status_code == 200 + assert response.json() is None diff --git a/tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py b/tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py index 7153fb909a8..c3bec658457 100644 --- a/tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_gateway_request_endpoints.py @@ -1,47 +1,25 @@ -import os from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Final -from unittest.mock import AsyncMock, MagicMock, call, patch - -from tests._master_key import MASTER_KEY - -# Patching ``litellm.proxy.proxy_server.prisma_client`` imports that module, whose -# module-level setup reads DATABASE_URL and LITELLM_MASTER_KEY. Tier-zero runners -# set neither, so pin throwaways first, as test_component_allowlists.py does. The -# prior values are restored below so a non-postgres URL cannot leak into sibling -# tests sharing the xdist worker and make them treat a phantom database as live. -_THROWAWAY_ENV = { - "DATABASE_URL": "sqlite:///:memory:", - "LITELLM_MASTER_KEY": MASTER_KEY, -} -_PRE_EXISTING_ENV = {key: os.environ.get(key) for key in _THROWAWAY_ENV} -for _key, _value in _THROWAWAY_ENV.items(): - os.environ.setdefault(_key, _value) +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient +from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import get_gateway_activity_repository from litellm.proxy.management_endpoints.gateway_request_endpoints import ( - _ACTIVITY_SQL, - _AggregateRow, _default_range, _fold_by_date, _fold_by_route, _fold_by_status_code, - _StatusCodeAggregateRow, get_gateway_daily_activity, router, ) - -for _key, _previous in _PRE_EXISTING_ENV.items(): - if _previous is None: - os.environ.pop(_key, None) - else: - os.environ[_key] = _previous +from litellm.repositories.gateway_activity_repository import GatewayActivityRepository # The handler stamps "today" from the wall clock, so any assertion that names a # date has to pin it. Recomputing the expected range in the assertion instead @@ -67,8 +45,8 @@ def _row( route: str = "/chat/completions", successful: int = 0, failed: int = 0, -) -> _AggregateRow: - return _AggregateRow( +) -> GatewayRouteActivity: + return GatewayRouteActivity( kind="route", date=date, category=category, @@ -86,9 +64,7 @@ def _prisma_returning(rows: Sequence[Mapping[str, str | int | None]] | None) -> client: Final = MagicMock() client.db = MagicMock() - async def query_raw(query: str, *_args: str) -> Sequence[Mapping[str, str | int | None]] | None: - if query != _ACTIVITY_SQL: - raise AssertionError(f"Unexpected gateway activity query: {query}") + async def query_raw(_query: str, *_args: str) -> Sequence[Mapping[str, str | int | None]] | None: return rows client.db.query_raw = AsyncMock(side_effect=query_raw) @@ -171,9 +147,9 @@ class TestFoldByStatusCode: def test_orders_by_failed_count_then_status_code(self): folded: Final = _fold_by_status_code( ( - _StatusCodeAggregateRow(kind="status", status_code=503, failed_requests=1), - _StatusCodeAggregateRow(kind="status", status_code=500, failed_requests=2), - _StatusCodeAggregateRow(kind="status", status_code=429, failed_requests=2), + GatewayStatusCodeActivity(kind="status", status_code=503, failed_requests=1), + GatewayStatusCodeActivity(kind="status", status_code=500, failed_requests=2), + GatewayStatusCodeActivity(kind="status", status_code=429, failed_requests=2), ) ) assert [(entry.status_code, entry.failed_requests) for entry in folded] == [(429, 2), (500, 2), (503, 1)] @@ -191,11 +167,11 @@ class TestGatewayDailyActivityEndpoint: ], ) async def test_refuses_every_non_admin_role(self, role): - with patch("litellm.proxy.proxy_server.prisma_client", _prisma_returning([])): - with pytest.raises(HTTPException) as exc: - await get_gateway_daily_activity( - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role), - ) + with pytest.raises(HTTPException) as exc: + await get_gateway_daily_activity( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role), + activity=None, + ) assert exc.value.status_code == 403 @pytest.mark.asyncio @@ -204,17 +180,18 @@ class TestGatewayDailyActivityEndpoint: [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], ) async def test_serves_both_admin_roles(self, role): - with patch("litellm.proxy.proxy_server.prisma_client", _prisma_returning([])): - response = await get_gateway_daily_activity( - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role), - ) + response = await get_gateway_daily_activity( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_role=role), + activity=GatewayActivityRepository(_prisma_returning([])), + start_date="2026-01-01", + end_date="2026-01-31", + ) assert response.total_successful_requests == 0 @pytest.mark.asyncio async def test_reports_db_not_connected_rather_than_crashing(self): - with patch("litellm.proxy.proxy_server.prisma_client", None): - with pytest.raises(HTTPException) as exc: - await get_gateway_daily_activity(user_api_key_dict=_admin()) + with pytest.raises(HTTPException) as exc: + await get_gateway_daily_activity(user_api_key_dict=_admin(), activity=None) assert exc.value.status_code == 500 @pytest.mark.asyncio @@ -267,8 +244,12 @@ class TestGatewayDailyActivityEndpoint: }, ) prisma: Final = _prisma_returning(rows) - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - response: Final = await get_gateway_daily_activity(user_api_key_dict=_admin()) + response: Final = await get_gateway_daily_activity( + user_api_key_dict=_admin(), + activity=GatewayActivityRepository(prisma), + start_date="2026-08-03", + end_date="2026-08-04", + ) assert response.total_successful_requests == 16 assert response.total_failed_requests == 4 @@ -277,13 +258,16 @@ class TestGatewayDailyActivityEndpoint: assert [(entry.status_code, entry.failed_requests) for entry in response.by_status_code] == [(429, 3), (500, 1)] assert [entry.date for entry in response.by_date] == ["2026-08-03", "2026-08-04"] assert [entry.route for entry in response.by_route] == ["/chat/completions", "/embeddings"] - assert prisma.db.query_raw.call_count == 1 @pytest.mark.asyncio async def test_a_null_result_set_is_not_an_error(self): client = _prisma_returning(None) - with patch("litellm.proxy.proxy_server.prisma_client", client): - response = await get_gateway_daily_activity(user_api_key_dict=_admin()) + response = await get_gateway_daily_activity( + user_api_key_dict=_admin(), + activity=GatewayActivityRepository(client), + start_date="2026-01-01", + end_date="2026-01-31", + ) assert response.total_successful_requests == 0 assert response.by_date == () assert response.by_route == () @@ -305,23 +289,23 @@ class TestGatewayDailyActivityRoute: app = FastAPI() app.include_router(router) app.dependency_overrides[user_api_key_auth] = _admin - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - response = TestClient(app).get( - "/gateway/daily/activity", - params={"start_date": "2026-01-01", "end_date": "2026-01-31"}, - ) + app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma) + response = TestClient(app).get( + "/gateway/daily/activity", + params={"start_date": "2026-01-01", "end_date": "2026-01-31"}, + ) assert response.status_code == 200 - assert prisma.db.query_raw.call_args_list == [call(_ACTIVITY_SQL, "2026-01-01", "2026-01-31")] + assert prisma.db.query_raw.call_args_list[0].args[1:] == ("2026-01-01", "2026-01-31") def test_omitted_dates_fall_back_to_the_default_window(self, frozen_clock): prisma = _prisma_returning([]) app = FastAPI() app.include_router(router) app.dependency_overrides[user_api_key_auth] = _admin - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - response = TestClient(app).get("/gateway/daily/activity") + app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma) + response = TestClient(app).get("/gateway/daily/activity") assert response.status_code == 200 - assert prisma.db.query_raw.call_args_list == [call(_ACTIVITY_SQL, *_FROZEN_RANGE)] + assert prisma.db.query_raw.call_args_list[0].args[1:] == _FROZEN_RANGE def test_serialized_response_carries_the_documented_shape(self): prisma: Final = _prisma_returning( @@ -349,8 +333,8 @@ class TestGatewayDailyActivityRoute: app: Final = FastAPI() app.include_router(router) app.dependency_overrides[user_api_key_auth] = _admin - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - body: Final = TestClient(app).get("/gateway/daily/activity").json() + app.dependency_overrides[get_gateway_activity_repository] = lambda: GatewayActivityRepository(prisma) + body: Final = TestClient(app).get("/gateway/daily/activity").json() assert body == { "total_successful_requests": 7, diff --git a/tests/unit/repositories/test_gateway_activity_repository.py b/tests/unit/repositories/test_gateway_activity_repository.py new file mode 100644 index 00000000000..e88702c5bdc --- /dev/null +++ b/tests/unit/repositories/test_gateway_activity_repository.py @@ -0,0 +1,47 @@ +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import ValidationError + +from litellm.models.gateway_activity import GatewayRouteActivity, GatewayStatusCodeActivity +from litellm.repositories.gateway_activity_repository import GatewayActivityRepository + + +@pytest.mark.asyncio +async def test_gateway_activity_decodes_both_row_kinds_with_nullable_union_columns() -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[ + { + "kind": "route", "date": "2026-01-02", "category": "llm", "route": "/chat/completions", + "successful_requests": 7, "failed_requests": 2, "status_code": None, + }, + { + "kind": "status", "date": None, "category": None, "route": None, + "successful_requests": 0, "failed_requests": 2, "status_code": 429, + }, + ]) + result: Final = await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31") + assert result == ( + GatewayRouteActivity( + kind="route", date="2026-01-02", category="llm", route="/chat/completions", + successful_requests=7, failed_requests=2, + ), + GatewayStatusCodeActivity(kind="status", status_code=429, failed_requests=2), + ) + + +@pytest.mark.asyncio +async def test_gateway_activity_rejects_malformed_database_rows() -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[{"kind": "status", "status_code": None, "failed_requests": 2}]) + with pytest.raises(ValidationError): + await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31") + + +@pytest.mark.asyncio +async def test_gateway_activity_does_not_hide_database_failure() -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(side_effect=RuntimeError("reader unavailable")) + with pytest.raises(RuntimeError, match="reader unavailable"): + await GatewayActivityRepository(prisma).get_activity("2026-01-01", "2026-01-31")