mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
780c504e00
* test: migrate tests/proxy_behavior into integration and unit tiers (partial) * test(integration): run settings-dependent management tests on an owned proxy * test(integration): assert stored values in the rewritten management rows * test(integration): pin sibling team allocations against the org guaranteed throughput cap * test(integration): keep Prisma clients out of the shared integration conftest import path --------- Co-authored-by: yuneng <yuneng@berri.ai>
35 lines
1 KiB
Python
35 lines
1 KiB
Python
import os
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from typing import Final
|
|
|
|
from prisma import Prisma
|
|
from prisma.types import DatasourceOverride
|
|
|
|
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
|
|
|
|
|
@asynccontextmanager
|
|
async def connected_prisma(database_url: str | None = None) -> AsyncIterator[Prisma]:
|
|
client: Final = Prisma(
|
|
datasource=DatasourceOverride(url=database_url)
|
|
) if database_url is not None else Prisma()
|
|
await client.connect()
|
|
try:
|
|
yield client
|
|
finally:
|
|
await client.disconnect()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def connected_proxy_prisma(database_url: str | None = None) -> AsyncIterator[PrismaClient]:
|
|
client: Final = PrismaClient(
|
|
database_url=database_url or os.environ["DATABASE_URL"],
|
|
proxy_logging_obj=ProxyLogging(UserApiKeyCache()),
|
|
)
|
|
await client.connect()
|
|
try:
|
|
yield client
|
|
finally:
|
|
await client.disconnect()
|