muhammadjalaludin
Implement service layers for backup, model management, file handling, tokenization, and training
1424118 | # 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) | |
| 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() | |