diff --git a/tests/integration/management/test_project_lifecycle.py b/tests/integration/management/test_project_lifecycle.py index 6982fbd581b..212229e3535 100644 --- a/tests/integration/management/test_project_lifecycle.py +++ b/tests/integration/management/test_project_lifecycle.py @@ -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", diff --git a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index e075d2c8aca..98bb2969201 100644 --- a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -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: diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index df285ab5a69..37d004b251b 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -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