scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

test(proxy): make ownership tests independent of runner salt and clock

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-10-02 22:52:42 +00:00
parent aba2597203
commit b5cfe66654
3 changed files with 32 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
tests/integration/management/test_project_lifecycle.py
View file

@ -13,6 +13,8 @@ from pydantic import JsonValue
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
_OWNED_PROXY_SALT_KEY: Final = "sk-integration-salt"
def _project_rows(project_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
@ -31,7 +33,14 @@ def _key_rows(key: str) -> list[dict[str, JsonValue]]:
)
def _cli_session_token(user_id: str, team_id: str, *, max_budget: float | None = None) -> str:
def _cli_session_token(
user_id: str,
team_id: str,
*,
monkeypatch: pytest.MonkeyPatch,
max_budget: float | None = None,
) -> str:
monkeypatch.setenv("LITELLM_SALT_KEY", _OWNED_PROXY_SALT_KEY)
user: Final = LiteLLM_UserTable(
user_id=user_id,
user_role="internal_user",
@ -57,7 +66,12 @@ def _discard_unexpected_key(candidate: Gateway, response: httpx.Response) -> Non
@pytest.fixture(scope="module")
def ownership_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
with gateway_from_environment() as gateway:
with owned_proxy(gateway, tmp_path_factory.mktemp("project-team-ownership"), {}, workers=2) as candidate:
with owned_proxy(
gateway,
tmp_path_factory.mktemp("project-team-ownership"),
{"LITELLM_SALT_KEY": _OWNED_PROXY_SALT_KEY},
workers=2,
) as candidate:
yield candidate
@ -347,14 +361,18 @@ def test_key_regenerate_routes_reject_missing_project_without_changing_key(owner
assert chat.status_code == 200, chat.text
def test_key_generate_budget_ceiling_precedes_project_ownership(ownership_gateway: Gateway) -> None:
def test_key_generate_budget_ceiling_precedes_project_ownership(
ownership_gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
with ownership_gateway.scenario() as scenario:
model: Final = scenario.model()
caller_team: Final = scenario.team(models=[model])
owner_team: Final = scenario.team(models=[model])
project: Final = scenario.project(owner_team, models=[model])
caller_id: Final = scenario.member(caller_team, role="admin")
caller_token: Final = _cli_session_token(caller_id, caller_team, max_budget=1)
caller_token: Final = _cli_session_token(
caller_id, caller_team, monkeypatch=monkeypatch, max_budget=1
)
response: Final = ownership_gateway.request(
"POST",
"/key/generate",
@ -379,7 +397,7 @@ def test_key_generate_budget_ceiling_precedes_project_ownership(ownership_gatewa
def test_key_regenerate_output_estimate_admin_error_precedes_project_ownership(
ownership_gateway: Gateway,
ownership_gateway: Gateway, monkeypatch: pytest.MonkeyPatch
) -> None:
with ownership_gateway.scenario() as scenario:
model: Final = scenario.model()
@ -388,7 +406,7 @@ def test_key_regenerate_output_estimate_admin_error_precedes_project_ownership(
project: Final = scenario.project(project_team, models=[model])
caller_id: Final = scenario.member(project_team, role="admin")
key: Final = scenario.key(team_id=project_team, user_id=caller_id, models=[model])
caller_token: Final = _cli_session_token(caller_id, project_team)
caller_token: Final = _cli_session_token(caller_id, project_team, monkeypatch=monkeypatch)
before: Final = _key_rows(key)
response: Final = ownership_gateway.request(
"POST",

8
tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py
View file

@ -1354,7 +1354,8 @@ async def test_update_project_slow_writer_ownership_reads_do_not_stall_db_tracke
async def slow_project_lookup(*, where: Mapping[str, object]) -> LiteLLM_ProjectTable:
assert where == {"project_id": project_id}
await asyncio.sleep(0.2)
await asyncio.sleep(0)
await asyncio.sleep(0)
return LiteLLM_ProjectTable(project_id=project_id, team_id=source_team_id)
async def slow_key_count(*, where: Mapping[str, object]) -> int:
@ -1362,12 +1363,13 @@ async def test_update_project_slow_writer_ownership_reads_do_not_stall_db_tracke
"project_id": project_id,
"OR": [{"team_id": {"not": destination_team_id}}, {"team_id": None}],
}
await asyncio.sleep(0.2)
await asyncio.sleep(0)
await asyncio.sleep(0)
return 0
mock_prisma.writer_db.litellm_projecttable.find_unique = mock.AsyncMock(side_effect=slow_project_lookup)
mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(side_effect=slow_key_count)
monkeypatch.setattr("litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.05)
monkeypatch.setattr("litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.0)
db_lookup_stall_tracker.clear()
try:

5
tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -20695,11 +20695,12 @@ async def test_key_generation_slow_writer_project_lookup_does_not_stall_db_track
) -> None:
user_api_key_cache: Final = await _cache_with_project(_OWNED_PROJECT, [], team_id=_OWNERSHIP_PROJECT_TEAM)
prisma_client: Final = _configure_key_endpoints(monkeypatch, user_api_key_cache)
monkeypatch.setattr("litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.05)
monkeypatch.setattr("litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.0)
async def slow_project_lookup(*, where: Mapping[str, object]) -> LiteLLM_ProjectTable:
assert where == {"project_id": _OWNED_PROJECT}
await asyncio.sleep(0.2)
await asyncio.sleep(0)
await asyncio.sleep(0)
return LiteLLM_ProjectTable(project_id=_OWNED_PROJECT, team_id=_OWNERSHIP_PROJECT_TEAM)
prisma_client.writer_db.litellm_projecttable.find_unique = slow_project_lookup