scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-28 01:32:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 95

test(proxy): type the shared test_connection helpers

Browse source
This commit is contained in:
Ishkirat-Singh 2026-09-09 18:25:44 +05:30
parent 36912cf03f
commit 03ee545366
1 changed files with 9 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -33,6 +33,7 @@ from litellm.proxy.health_endpoints._health_endpoints import (
from litellm.proxy.health_endpoints._health_endpoints import (
test_model_connection as health_test_model_connection,
)
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
# Import shared proxy test helpers from conftest
from tests.test_litellm.proxy.conftest import create_proxy_test_client
@ -4181,9 +4182,7 @@ async def test_health_services_endpoint_pointfive_blocks_non_admin(monkeypatch,
logger_class.assert_not_called()
def _configured_non_team_deployment(mode: str | None = "chat"):
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
def _configured_non_team_deployment(mode: str | None = "chat") -> Deployment:
return Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(
@ -4195,7 +4194,7 @@ def _configured_non_team_deployment(mode: str | None = "chat"):
)
def _internal_user():
def _internal_user() -> UserAPIKeyAuth:
return UserAPIKeyAuth(
token="internal-user-token",
user_id="internal-user",
@ -4205,7 +4204,7 @@ def _internal_user():
@contextmanager
def _probe_environment(
deployment,
deployment: Deployment,
*,
key_check: AsyncMock | None = None,
user_lookup: AsyncMock | None = None,
@ -4257,7 +4256,9 @@ def _probe_environment(
yield
async def _probe_as_internal_user(litellm_params, mode="chat", deployment_id="non-team-deployment-id"):
async def _probe_as_internal_user(
litellm_params: dict[str, object], mode: str | None = "chat", deployment_id: str = "non-team-deployment-id"
) -> dict[str, object]:
return await health_test_model_connection(
request=MagicMock(),
mode=mode,
@ -4327,7 +4328,7 @@ async def test_test_model_connection_denies_internal_user_without_model_access()
],
ids=["api_base", "model", "mode"],
)
async def test_test_model_connection_keeps_overrides_admin_only(litellm_params, mode):
async def test_test_model_connection_keeps_overrides_admin_only(litellm_params: dict[str, object], mode: str) -> None:
from fastapi import HTTPException
key_check = AsyncMock(return_value=True)
@ -4349,7 +4350,6 @@ async def test_test_model_connection_keeps_team_deployments_admin_only_for_non_a
from fastapi import HTTPException
from litellm.proxy._types import LiteLLM_TeamTable
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
team_deployment = Deployment(
model_name="gpt-4o",
@ -4399,8 +4399,7 @@ async def test_test_model_connection_tolerates_missing_user_record_for_non_admin
@pytest.mark.asyncio
async def test_test_model_connection_accepts_mode_the_probe_would_infer():
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
async def test_test_model_connection_accepts_mode_the_probe_would_infer() -> None:
deployment = Deployment(
model_name="gpt-4o",