muhammadjalaludin
Implement service layers for backup, model management, file handling, tokenization, and training
1424118 | # ta_ml/tests/api/test_predict.py | |
| """ | |
| Integration tests for the /predict API endpoint. | |
| """ | |
| from fastapi.testclient import TestClient | |
| from unittest.mock import MagicMock | |
| import pytest | |
| from main import app | |
| from services.prediction_service import PredictionService | |
| from api.dependencies import get_prediction_service | |
| # The TestClient allows us to make HTTP requests to our FastAPI app in tests | |
| client = TestClient(app) | |
| def mock_prediction_service(): | |
| """Mocks the prediction service to control its behavior during API tests.""" | |
| mock = MagicMock(spec=PredictionService) | |
| # The 'spec' argument ensures the mock has the same methods as the real class | |
| return mock | |
| def test_predict_success(mock_prediction_service): | |
| """Tests the happy path for the /predict endpoint.""" | |
| # Arrange | |
| expected_response = {"predicted_label": "Test Label", "confidence": 0.99} | |
| mock_prediction_service.predict.return_value = expected_response | |
| # Override the real dependency with our mock for this test | |
| app.dependency_overrides[get_prediction_service] = lambda: mock_prediction_service | |
| # Act | |
| response = client.post("/api/predict", json={"text": "some transaction"}) | |
| # Assert | |
| assert response.status_code == 200 | |
| assert response.json() == expected_response | |
| mock_prediction_service.predict.assert_called_once_with("some transaction") | |
| # Clean up the override after the test | |
| app.dependency_overrides.clear() | |
| def test_predict_invalid_input(mock_prediction_service): | |
| """Tests the endpoint's response when the service raises a ValueError.""" | |
| # Arrange | |
| mock_prediction_service.predict.side_effect = ValueError( | |
| "Input cannot be empty") | |
| app.dependency_overrides[get_prediction_service] = lambda: mock_prediction_service | |
| # Act | |
| response = client.post("/api/predict", json={"text": ""}) # Text is empty | |
| # Assert | |
| assert response.status_code == 400 | |
| assert response.json() == {"detail": "Input cannot be empty"} | |
| app.dependency_overrides.clear() | |
| def test_predict_missing_text_field(): | |
| """Tests the validation handled by FastAPI/Pydantic for a malformed request.""" | |
| # Act | |
| # We send a request without the required 'text' field | |
| response = client.post("/api/predict", json={"wrong_field": "some data"}) | |
| # Assert | |
| # FastAPI should automatically return a 422 Unprocessable Entity error | |
| assert response.status_code == 422 | |
| response_json = response.json() | |
| assert "detail" in response_json | |
| assert response_json["detail"][0]["msg"] == "Field required" | |
| assert response_json["detail"][0]["loc"] == ["body", "text"] | |
| def test_predict_internal_server_error(mock_prediction_service): | |
| """Tests the endpoint's generic 500 error handling.""" | |
| # Arrange | |
| # Simulate an unexpected error in the service | |
| mock_prediction_service.predict.side_effect = Exception( | |
| "A wild error appeared!") | |
| app.dependency_overrides[get_prediction_service] = lambda: mock_prediction_service | |
| # Act | |
| response = client.post("/api/predict", json={"text": "this will fail"}) | |
| # Assert | |
| assert response.status_code == 500 | |
| assert response.json() == { | |
| "detail": "An internal error occurred during prediction."} | |
| app.dependency_overrides.clear() | |