162 lines
4.5 KiB
Python
162 lines
4.5 KiB
Python
"""Pytest fixtures for LLM inference tests."""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
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
|
|
from llm_inference.schemas import CompletionChunk, CompletionResponse
|
|
from llm_inference.types import (
|
|
BackendType,
|
|
ChatMessage,
|
|
Choice,
|
|
Delta,
|
|
StreamChoice,
|
|
Usage,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_settings_cache() -> None:
|
|
"""Clear settings cache before each test."""
|
|
SettingsCache.clear()
|
|
|
|
|
|
@pytest.fixture
|
|
def test_settings() -> LLMSettings:
|
|
"""Create test settings with mocked values."""
|
|
return LLMSettings(
|
|
default_backend="litellm",
|
|
default_model="test-model",
|
|
openrouter_api_key="test-openrouter-key",
|
|
openai_api_key="test-openai-key",
|
|
enable_vllm=False,
|
|
enable_llamacpp=False,
|
|
host="127.0.0.1",
|
|
port=8100,
|
|
external_url="http://localhost:8100",
|
|
api_tokens=None, # Auth disabled by default in tests
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_completion_response() -> CompletionResponse:
|
|
"""Create a mock completion response."""
|
|
return CompletionResponse(
|
|
id="test-completion-id",
|
|
created=1234567890,
|
|
model="test-model",
|
|
choices=[
|
|
Choice(
|
|
index=0,
|
|
message=ChatMessage(role="assistant", content="Hello! How can I help?"),
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
|
backend="litellm",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_completion_chunks() -> list[CompletionChunk]:
|
|
"""Create mock streaming completion chunks."""
|
|
return [
|
|
CompletionChunk(
|
|
id="test-chunk-1",
|
|
created=1234567890,
|
|
model="test-model",
|
|
choices=[
|
|
StreamChoice(
|
|
index=0,
|
|
delta=Delta(role="assistant", content="Hello"),
|
|
finish_reason=None,
|
|
)
|
|
],
|
|
),
|
|
CompletionChunk(
|
|
id="test-chunk-2",
|
|
created=1234567891,
|
|
model="test-model",
|
|
choices=[
|
|
StreamChoice(
|
|
index=0,
|
|
delta=Delta(content="!"),
|
|
finish_reason=None,
|
|
)
|
|
],
|
|
),
|
|
CompletionChunk(
|
|
id="test-chunk-3",
|
|
created=1234567892,
|
|
model="test-model",
|
|
choices=[
|
|
StreamChoice(
|
|
index=0,
|
|
delta=Delta(),
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_litellm_response() -> MagicMock:
|
|
"""Create a mock LiteLLM response object."""
|
|
response = MagicMock()
|
|
response.id = "test-completion-id"
|
|
response.created = 1234567890
|
|
response.model = "test-model"
|
|
|
|
choice = MagicMock()
|
|
choice.index = 0
|
|
choice.message.role = "assistant"
|
|
choice.message.content = "Hello! How can I help?"
|
|
choice.finish_reason = "stop"
|
|
response.choices = [choice]
|
|
|
|
response.usage.prompt_tokens = 10
|
|
response.usage.completion_tokens = 5
|
|
response.usage.total_tokens = 15
|
|
|
|
return response
|
|
|
|
|
|
@pytest.fixture
|
|
def client_with_mock_backend(
|
|
test_settings: LLMSettings, mock_completion_response: CompletionResponse
|
|
) -> LLMClient:
|
|
"""Create an LLMClient with mocked backend."""
|
|
# Inject test settings into cache
|
|
SettingsCache.set(test_settings)
|
|
|
|
client = LLMClient(settings=test_settings)
|
|
|
|
# Mock the backend's complete method
|
|
backend = client.registry.get(BackendType.LITELLM)
|
|
backend.complete = AsyncMock(return_value=mock_completion_response)
|
|
|
|
return client
|
|
|
|
|
|
@pytest.fixture
|
|
def app_client(test_settings: LLMSettings) -> TestClient:
|
|
"""Create a FastAPI TestClient with test settings."""
|
|
# Inject test settings into cache BEFORE creating the app
|
|
SettingsCache.set(test_settings)
|
|
|
|
# Initialize concurrency limiter for tests
|
|
init_concurrency_limiter(test_settings.max_concurrent_completions)
|
|
|
|
app = create_app()
|
|
|
|
# Also set in app.state for handlers that access it
|
|
app.state.settings = test_settings
|
|
app.state.client = LLMClient(settings=test_settings)
|
|
|
|
return TestClient(app)
|