228 lines
6.7 KiB
Python
228 lines
6.7 KiB
Python
"""Tests for API routes."""
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from src.engine.backend_manager import BackendManager
|
|
from src.main import app
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_backend_manager():
|
|
"""Create a mock backend manager."""
|
|
manager = MagicMock(spec=BackendManager)
|
|
manager.backend_type = "ollama"
|
|
|
|
# Mock health_check
|
|
manager.health_check = AsyncMock(return_value=True)
|
|
|
|
# Mock list_models
|
|
manager.list_models = AsyncMock(return_value=[
|
|
{
|
|
"id": "tinyllama:latest",
|
|
"object": "model",
|
|
"created": 1234567890,
|
|
"owned_by": "ollama",
|
|
"permission": [],
|
|
"root": "tinyllama:latest",
|
|
"parent": None,
|
|
}
|
|
])
|
|
|
|
# Mock chat_completion
|
|
manager.chat_completion = AsyncMock(return_value={
|
|
"id": "chatcmpl-test",
|
|
"created": 1234567890,
|
|
"content": "Hello! How can I help you?",
|
|
"finish_reason": "stop",
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 8,
|
|
"total_tokens": 18,
|
|
},
|
|
})
|
|
|
|
return manager
|
|
|
|
|
|
@pytest.fixture
|
|
def client(mock_backend_manager):
|
|
"""Create test client with mocked backend."""
|
|
app.state.backend_manager = mock_backend_manager
|
|
return TestClient(app)
|
|
|
|
|
|
class TestHealthEndpoints:
|
|
"""Test health check endpoints."""
|
|
|
|
def test_health_check(self, client, mock_backend_manager):
|
|
"""Test main health endpoint."""
|
|
response = client.get("/health")
|
|
assert response.status_code == 200
|
|
|
|
data = response.json()
|
|
assert "status" in data
|
|
assert "timestamp" in data
|
|
assert "version" in data
|
|
assert "dependencies" in data
|
|
|
|
def test_liveness_check(self, client):
|
|
"""Test liveness endpoint."""
|
|
response = client.get("/health/live")
|
|
assert response.status_code == 200
|
|
|
|
data = response.json()
|
|
assert data["status"] == "alive"
|
|
assert "timestamp" in data
|
|
|
|
def test_readiness_check(self, client, mock_backend_manager):
|
|
"""Test readiness endpoint."""
|
|
response = client.get("/health/ready")
|
|
assert response.status_code == 200
|
|
|
|
data = response.json()
|
|
assert "ready" in data
|
|
assert "checks" in data
|
|
assert "timestamp" in data
|
|
|
|
|
|
class TestModelsEndpoint:
|
|
"""Test models listing endpoint."""
|
|
|
|
def test_list_models(self, client, mock_backend_manager):
|
|
"""Test listing models."""
|
|
response = client.get("/v1/models")
|
|
assert response.status_code == 200
|
|
|
|
data = response.json()
|
|
assert data["object"] == "list"
|
|
assert "data" in data
|
|
assert len(data["data"]) > 0
|
|
|
|
def test_list_models_structure(self, client, mock_backend_manager):
|
|
"""Test model structure matches OpenAI format."""
|
|
response = client.get("/v1/models")
|
|
data = response.json()
|
|
|
|
model = data["data"][0]
|
|
assert "id" in model
|
|
assert "object" in model
|
|
assert model["object"] == "model"
|
|
|
|
|
|
class TestChatCompletionEndpoint:
|
|
"""Test chat completion endpoint."""
|
|
|
|
def test_chat_completion_basic(self, client, mock_backend_manager):
|
|
"""Test basic chat completion."""
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": "tinyllama",
|
|
"messages": [
|
|
{"role": "user", "content": "Hello!"}
|
|
],
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
data = response.json()
|
|
assert "id" in data
|
|
assert "choices" in data
|
|
assert "usage" in data
|
|
assert data["object"] == "chat.completion"
|
|
|
|
def test_chat_completion_with_options(self, client, mock_backend_manager):
|
|
"""Test chat completion with all options."""
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": "tinyllama",
|
|
"messages": [
|
|
{"role": "system", "content": "You are helpful."},
|
|
{"role": "user", "content": "Hello!"},
|
|
],
|
|
"max_tokens": 100,
|
|
"temperature": 0.5,
|
|
"top_p": 0.9,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
def test_chat_completion_empty_messages_rejected(self, client):
|
|
"""Test empty messages are rejected."""
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": "tinyllama",
|
|
"messages": [],
|
|
},
|
|
)
|
|
assert response.status_code == 422 # Validation error
|
|
|
|
def test_chat_completion_invalid_role_rejected(self, client):
|
|
"""Test invalid role is rejected."""
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": "tinyllama",
|
|
"messages": [
|
|
{"role": "invalid", "content": "Hello!"}
|
|
],
|
|
},
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
def test_chat_completion_invalid_temperature_rejected(self, client):
|
|
"""Test invalid temperature is rejected."""
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": "tinyllama",
|
|
"messages": [
|
|
{"role": "user", "content": "Hello!"}
|
|
],
|
|
"temperature": 5.0, # Too high
|
|
},
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
def test_chat_completion_response_structure(self, client, mock_backend_manager):
|
|
"""Test response structure matches OpenAI format."""
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": "tinyllama",
|
|
"messages": [
|
|
{"role": "user", "content": "Hello!"}
|
|
],
|
|
},
|
|
)
|
|
data = response.json()
|
|
|
|
# Check structure
|
|
assert "id" in data
|
|
assert "object" in data
|
|
assert "created" in data
|
|
assert "model" in data
|
|
assert "choices" in data
|
|
assert "usage" in data
|
|
|
|
# Check choices structure
|
|
choice = data["choices"][0]
|
|
assert "index" in choice
|
|
assert "message" in choice
|
|
assert "finish_reason" in choice
|
|
|
|
# Check message structure
|
|
message = choice["message"]
|
|
assert "role" in message
|
|
assert "content" in message
|
|
|
|
# Check usage structure
|
|
usage = data["usage"]
|
|
assert "prompt_tokens" in usage
|
|
assert "completion_tokens" in usage
|
|
assert "total_tokens" in usage
|