mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
7a38fd33c1
Removes the suite README — its contents (local repro, layout, conventions) were either restated by the file structure or already covered by the workflow YAML and pyproject.toml. Trims docstrings and inline comments across every test file to keep only non-obvious WHY (the masking ``_get_user_in_team`` reads, the LiteLLM_VerificationTokenView models-can't- be-NULL gotcha, the org_admin/peer-visibility surprise, the rotation contract). Suite still 129 green locally.
128 lines
4 KiB
Python
128 lines
4 KiB
Python
"""Session-scoped async ASGI client for HTTP-boundary behavior tests."""
|
|
|
|
import os
|
|
import tempfile
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from typing import AsyncIterator
|
|
|
|
import httpx
|
|
import pytest_asyncio
|
|
import yaml
|
|
|
|
|
|
MASTER_KEY = "sk-1234"
|
|
SCRATCH_PREFIX = "scratch-"
|
|
|
|
|
|
def _write_minimal_proxy_config() -> str:
|
|
config = {
|
|
"general_settings": {"master_key": MASTER_KEY},
|
|
"litellm_settings": {},
|
|
}
|
|
database_url = os.environ.get("DATABASE_URL")
|
|
if database_url:
|
|
config["general_settings"]["database_url"] = database_url
|
|
f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
|
|
yaml.dump(config, f)
|
|
f.close()
|
|
return f.name
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="session")
|
|
async def proxy_app():
|
|
from litellm.proxy import proxy_server
|
|
from litellm.proxy.proxy_server import (
|
|
app,
|
|
cleanup_router_config_variables,
|
|
initialize,
|
|
proxy_startup_event,
|
|
)
|
|
|
|
cleanup_router_config_variables()
|
|
config_path = _write_minimal_proxy_config()
|
|
|
|
# proxy_startup_event re-reads master_key from LITELLM_MASTER_KEY and
|
|
# unconditionally overwrites the global, even when initialize() already
|
|
# set it from the config YAML. Without these env vars, the entire auth
|
|
# stack degrades to user_id=None / non-PROXY_ADMIN for every token.
|
|
os.environ.setdefault("LITELLM_MASTER_KEY", MASTER_KEY)
|
|
os.environ.setdefault("CONFIG_FILE_PATH", config_path)
|
|
|
|
await initialize(config=config_path)
|
|
|
|
# /key/regenerate is gated behind premium_user; flipping it lets the matrix
|
|
# pin authz behavior instead of the licensing gate.
|
|
proxy_server.premium_user = True
|
|
|
|
async with proxy_startup_event(app):
|
|
proxy_server.premium_user = True # lifespan re-runs _license_check
|
|
# The lifespan fires check_view_exists() as a background task; on a
|
|
# fresh DB the first auth call races it and resolves user_id=None.
|
|
if proxy_server.prisma_client is not None:
|
|
await proxy_server.prisma_client.check_view_exists()
|
|
yield app
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="session")
|
|
async def proxy_client(proxy_app) -> AsyncIterator[httpx.AsyncClient]:
|
|
transport = httpx.ASGITransport(app=proxy_app)
|
|
async with httpx.AsyncClient(
|
|
transport=transport, base_url="http://testserver"
|
|
) as client:
|
|
yield client
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="session")
|
|
async def prisma(proxy_app):
|
|
from litellm.proxy import proxy_server
|
|
|
|
assert proxy_server.prisma_client is not None
|
|
return proxy_server.prisma_client
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="session")
|
|
async def world(prisma):
|
|
from .actors import seed_world
|
|
|
|
return await seed_world(prisma)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Scratch:
|
|
prefix: str
|
|
|
|
def tag(self, suffix: str = "") -> str:
|
|
return f"{self.prefix}-{suffix}" if suffix else self.prefix
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def scratch(prisma):
|
|
handle = Scratch(prefix=f"{SCRATCH_PREFIX}{uuid.uuid4().hex[:12]}")
|
|
try:
|
|
yield handle
|
|
finally:
|
|
# Children before parents to avoid FK violations.
|
|
await prisma.db.litellm_verificationtoken.delete_many(
|
|
where={
|
|
"OR": [
|
|
{"key_alias": {"startswith": handle.prefix}},
|
|
{"key_name": {"startswith": handle.prefix}},
|
|
]
|
|
}
|
|
)
|
|
await prisma.db.litellm_teammembership.delete_many(
|
|
where={"team_id": {"startswith": handle.prefix}}
|
|
)
|
|
await prisma.db.litellm_organizationmembership.delete_many(
|
|
where={"user_id": {"startswith": handle.prefix}}
|
|
)
|
|
await prisma.db.litellm_teamtable.delete_many(
|
|
where={"team_id": {"startswith": handle.prefix}}
|
|
)
|
|
await prisma.db.litellm_usertable.delete_many(
|
|
where={"user_id": {"startswith": handle.prefix}}
|
|
)
|
|
await prisma.db.litellm_budgettable.delete_many(
|
|
where={"budget_id": {"startswith": handle.prefix}}
|
|
)
|