test: cover predictor routes

Author: arthurhenrique

{{cookiecutter.project_slug}}/app/api/routes/predictor.py

@@ -1,8 +1,10 @@
 import json
+from pathlib import Path
 
 import joblib
 from core.config import INPUT_EXAMPLE
 from fastapi import APIRouter, HTTPException
+from fastapi.concurrency import run_in_threadpool
 from models.prediction import (
     HealthResponse,
     MachineLearningDataInput,
@@ -33,11 +35,14 @@ async def predict(data_input: MachineLearningDataInput):
         raise HTTPException(status_code=404, detail="'data_input' argument invalid!")
     try:
         data_point = data_input.get_np_array()
-        prediction = get_prediction(data_point)
+        prediction = await run_in_threadpool(get_prediction, data_point)
+        try:
+            prediction = float(prediction[0])
+        except (TypeError, IndexError, KeyError):
+            prediction = float(prediction)
         prediction_label = get_prediction_label(prediction)
-
     except Exception as err:
-        raise HTTPException(status_code=500, detail=f"Exception: {err}")
+        raise HTTPException(status_code=500, detail=f"Exception: {err}") from err
 
     return MachineLearningResponse(
         prediction=prediction, prediction_label=prediction_label
@@ -50,14 +55,11 @@ async def predict(data_input: MachineLearningDataInput):
     name="health:get-data",
 )
 async def health():
-    is_health = False
     try:
-        test_input = MachineLearningDataInput(
-            **json.loads(open(INPUT_EXAMPLE, "r").read())
-        )
+        content = await run_in_threadpool(Path(INPUT_EXAMPLE).read_text)
+        test_input = MachineLearningDataInput(**json.loads(content))
         test_point = test_input.get_np_array()
-        get_prediction(test_point)
-        is_health = True
-        return HealthResponse(status=is_health)
+        await run_in_threadpool(get_prediction, test_point)
+        return HealthResponse(status=True)
     except Exception:
         raise HTTPException(status_code=404, detail="Unhealthy")

{{cookiecutter.project_slug}}/app/models/prediction.py

@@ -5,6 +5,7 @@
 
 class MachineLearningResponse(BaseModel):
     prediction: float
+    prediction_label: str
 
 
 class HealthResponse(BaseModel):

{{cookiecutter.project_slug}}/tests/test_predictor.py

@@ -0,0 +1,38 @@
+import json
+import sys
+from pathlib import Path
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.append(str(Path(__file__).resolve().parents[1] / "app"))
+
+from api.routes import predictor
+from main import app
+
+
+client = TestClient(app)
+
+
+def sample_input():
+    example_path = Path(__file__).resolve().parents[1] / "ml" / "model" / "examples" / "example.json"
+    return json.loads(example_path.read_text())
+
+
+def test_predict_endpoint(monkeypatch):
+    monkeypatch.setattr(predictor, "get_prediction", lambda data_point: [1.0])
+    response = client.post("/api/v1/predict", json=sample_input())
+    assert response.status_code == 200
+    assert response.json() == {"prediction": 1.0, "prediction_label": "label ok"}
+
+
+def test_health_endpoint(monkeypatch):
+    monkeypatch.setattr(predictor, "get_prediction", lambda data_point: [1.0])
+    monkeypatch.setattr(
+        predictor,
+        "INPUT_EXAMPLE",
+        str(Path(__file__).resolve().parents[1] / "ml" / "model" / "examples" / "example.json"),
+    )
+    response = client.get("/api/v1/health")
+    assert response.status_code == 200
+    assert response.json() == {"status": True}