didi-lot1-ai/ai_platform/modules/llm-inference/tests/test_auth.py

224 lines
8.8 KiB
Python

"""Tests for API authentication."""
import pytest
from fastapi.testclient import TestClient
from llm_inference.api.app import create_app
from llm_inference.api.dependencies import init_concurrency_limiter
from llm_inference.client import LLMClient
from llm_inference.config import LLMSettings, SettingsCache
class TestAuthDisabled:
"""Tests when authentication is disabled (no tokens configured)."""
@pytest.fixture
def app_client_no_auth(self, test_settings: LLMSettings) -> TestClient:
"""Create test client with auth disabled."""
# test_settings has api_tokens=None by default
SettingsCache.set(test_settings)
init_concurrency_limiter(test_settings.max_concurrent_completions)
app = create_app()
app.state.settings = test_settings
app.state.client = LLMClient(settings=test_settings)
return TestClient(app)
def test_models_accessible_without_token(
self, app_client_no_auth: TestClient
) -> None:
"""Test that models endpoint works without token when auth disabled."""
response = app_client_no_auth.get("/v1/models")
assert response.status_code != 401
def test_backends_accessible_without_token(
self, app_client_no_auth: TestClient
) -> None:
"""Test that backends endpoint works without token when auth disabled."""
response = app_client_no_auth.get("/v1/backends")
assert response.status_code == 200
def test_health_accessible(self, app_client_no_auth: TestClient) -> None:
"""Test health endpoint accessible."""
response = app_client_no_auth.get("/health")
assert response.status_code == 200
def test_ready_accessible(self, app_client_no_auth: TestClient) -> None:
"""Test ready endpoint accessible."""
response = app_client_no_auth.get("/ready")
# May return 200 or 503 depending on backend health, but not 401
assert response.status_code != 401
class TestAuthEnabled:
"""Tests when authentication is enabled."""
@pytest.fixture
def auth_settings(self, test_settings: LLMSettings) -> LLMSettings:
"""Create settings with auth enabled."""
return LLMSettings(
default_backend=test_settings.default_backend,
enable_vllm=test_settings.enable_vllm,
enable_llamacpp=test_settings.enable_llamacpp,
external_url=test_settings.external_url,
openrouter_api_key=test_settings.openrouter_api_key,
api_tokens=frozenset(["test-token-1", "test-token-2"]),
)
@pytest.fixture
def app_client_auth(self, auth_settings: LLMSettings) -> TestClient:
"""Create test client with auth enabled."""
SettingsCache.set(auth_settings)
init_concurrency_limiter(auth_settings.max_concurrent_completions)
app = create_app()
app.state.settings = auth_settings
app.state.client = LLMClient(settings=auth_settings)
return TestClient(app)
def test_missing_token_returns_401(self, app_client_auth: TestClient) -> None:
"""Test 401 returned when token missing."""
response = app_client_auth.get("/v1/models")
assert response.status_code == 401
assert "WWW-Authenticate" in response.headers
assert response.headers["WWW-Authenticate"] == "Bearer"
data = response.json()
assert data["detail"]["error"] == "Authentication required"
def test_invalid_token_returns_401(self, app_client_auth: TestClient) -> None:
"""Test 401 returned for invalid token."""
response = app_client_auth.get(
"/v1/models",
headers={"Authorization": "Bearer invalid-token"},
)
assert response.status_code == 401
data = response.json()
assert data["detail"]["error"] == "Authentication failed"
def test_valid_token_allows_access(self, app_client_auth: TestClient) -> None:
"""Test valid token grants access."""
response = app_client_auth.get(
"/v1/models",
headers={"Authorization": "Bearer test-token-1"},
)
assert response.status_code == 200
def test_second_valid_token_works(self, app_client_auth: TestClient) -> None:
"""Test second token also works."""
response = app_client_auth.get(
"/v1/models",
headers={"Authorization": "Bearer test-token-2"},
)
assert response.status_code == 200
def test_malformed_header_missing_bearer(self, app_client_auth: TestClient) -> None:
"""Test malformed Authorization header without Bearer prefix."""
response = app_client_auth.get(
"/v1/models",
headers={"Authorization": "test-token-1"},
)
assert response.status_code == 401
def test_malformed_header_wrong_prefix(self, app_client_auth: TestClient) -> None:
"""Test Authorization header with wrong prefix."""
response = app_client_auth.get(
"/v1/models",
headers={"Authorization": "Basic test-token-1"},
)
assert response.status_code == 401
def test_health_excluded_from_auth(self, app_client_auth: TestClient) -> None:
"""Test /health accessible without token even when auth enabled."""
response = app_client_auth.get("/health")
assert response.status_code == 200
def test_ready_excluded_from_auth(self, app_client_auth: TestClient) -> None:
"""Test /ready accessible without token even when auth enabled."""
response = app_client_auth.get("/ready")
# May return 200 or 503 depending on backend health, but not 401
assert response.status_code != 401
def test_backends_requires_auth(self, app_client_auth: TestClient) -> None:
"""Test /v1/backends requires authentication."""
response = app_client_auth.get("/v1/backends")
assert response.status_code == 401
def test_backends_with_valid_token(self, app_client_auth: TestClient) -> None:
"""Test /v1/backends works with valid token."""
response = app_client_auth.get(
"/v1/backends",
headers={"Authorization": "Bearer test-token-1"},
)
assert response.status_code == 200
class TestTokenParsing:
"""Tests for token configuration parsing."""
def test_comma_separated_tokens(self) -> None:
"""Test parsing comma-separated tokens."""
settings = LLMSettings(
default_backend="litellm",
enable_vllm=False,
enable_llamacpp=False,
external_url="http://localhost:8100",
api_tokens="token1,token2,token3",
)
assert settings.api_tokens == frozenset(["token1", "token2", "token3"])
assert settings.auth_enabled is True
def test_empty_string_disables_auth(self) -> None:
"""Test empty string results in auth disabled."""
settings = LLMSettings(
default_backend="litellm",
enable_vllm=False,
enable_llamacpp=False,
external_url="http://localhost:8100",
api_tokens="",
)
assert settings.api_tokens is None
assert settings.auth_enabled is False
def test_whitespace_tokens_stripped(self) -> None:
"""Test whitespace around tokens is stripped."""
settings = LLMSettings(
default_backend="litellm",
enable_vllm=False,
enable_llamacpp=False,
external_url="http://localhost:8100",
api_tokens=" token1 , token2 , token3 ",
)
assert settings.api_tokens == frozenset(["token1", "token2", "token3"])
def test_none_disables_auth(self) -> None:
"""Test None results in auth disabled."""
settings = LLMSettings(
default_backend="litellm",
enable_vllm=False,
enable_llamacpp=False,
external_url="http://localhost:8100",
)
assert settings.api_tokens is None
assert settings.auth_enabled is False
def test_single_token(self) -> None:
"""Test single token without comma."""
settings = LLMSettings(
default_backend="litellm",
enable_vllm=False,
enable_llamacpp=False,
external_url="http://localhost:8100",
api_tokens="single-token",
)
assert settings.api_tokens == frozenset(["single-token"])
assert settings.auth_enabled is True
def test_whitespace_only_disables_auth(self) -> None:
"""Test whitespace-only string results in auth disabled."""
settings = LLMSettings(
default_backend="litellm",
enable_vllm=False,
enable_llamacpp=False,
external_url="http://localhost:8100",
api_tokens=" , , ",
)
assert settings.api_tokens is None
assert settings.auth_enabled is False