mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
ee2cb0b6e0
Database now holds the backend handle and its close operation only, and each feature provider builds its own repository from it, so _db never grows a field per feature. get_database checks the concrete type instead of validating a runtime Protocol on every request, and the banner audit runs as a FastAPI background task so the in-flight drain waits for it before the database closes The proxy-dependencies skill moves to the repo root .agents/skills and gains bridge rules against passing SQL text or identifiers from Python Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
36 lines
1 KiB
Python
36 lines
1 KiB
Python
import asyncio
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from litellm import _db
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError])
|
|
async def test_lifespan_closes_database_when_owner_fails(failure):
|
|
prisma: Final = MagicMock()
|
|
prisma.disconnect = AsyncMock()
|
|
database: Final = _db.Database(prisma=prisma)
|
|
|
|
async def fail():
|
|
async with _db.lifespan(database) as current:
|
|
assert current is database
|
|
prisma.disconnect.assert_not_awaited()
|
|
raise failure("owner failed")
|
|
|
|
with pytest.raises(failure):
|
|
await fail()
|
|
prisma.disconnect.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lifespan_without_database_preserves_owner_error():
|
|
async def fail():
|
|
async with _db.lifespan(None) as database:
|
|
assert database is None
|
|
raise RuntimeError("owner failed")
|
|
|
|
with pytest.raises(RuntimeError, match="owner failed"):
|
|
await fail()
|