File size: 1,923 Bytes
1424118 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | # ta_ml/tests/api/test_admin.py
"""
Integration tests for the /admin API endpoints.
"""
from fastapi.testclient import TestClient
from unittest.mock import MagicMock
import pytest
from main import app
from services.model_manager import ModelManager
from api.dependencies import get_model_manager
client = TestClient(app)
@pytest.fixture
def mock_model_manager():
"""Mocks the ModelManager service."""
return MagicMock(spec=ModelManager)
def test_reload_models_success(mock_model_manager):
"""Tests the happy path for the model reload endpoint."""
# Arrange
# The real ModelManager is complex to initialize, so we mock it.
# We also need to mock the services it creates.
app.state.model_manager = mock_model_manager
app.state.prediction_service = MagicMock()
app.state.tokenization_service = MagicMock()
app.state.training_service = MagicMock()
# Override the dependency to ensure our mock is used
app.dependency_overrides[get_model_manager] = lambda: mock_model_manager
# Act
response = client.post("/api/admin/reload-models")
# Assert
assert response.status_code == 200
assert response.json()["status"] == "Models reloaded successfully"
mock_model_manager.reload_models.assert_called_once()
app.dependency_overrides.clear()
def test_reload_models_failure(mock_model_manager):
"""Tests the endpoint's response when the reload process fails."""
# Arrange
mock_model_manager.reload_models.side_effect = RuntimeError(
"Could not download from Hub")
app.dependency_overrides[get_model_manager] = lambda: mock_model_manager
# Act
response = client.post("/api/admin/reload-models")
# Assert
assert response.status_code == 500
assert "Error reloading models" in response.json()["detail"]
assert "Could not download from Hub" in response.json()["detail"]
app.dependency_overrides.clear()
|