Add test cases for API endpoints including health check and spam prediction

tests/test_api.py

@@ -0,0 +1,36 @@
+# tests/test_api.py
+import pytest
+from fastapi.testclient import TestClient
+from main import app
+
+# --- FIX: Create a fixture that manages the startup/shutdown ---
+@pytest.fixture(scope="module")
+def client():
+    # The 'with' statement triggers the lifespan (startup) event
+    with TestClient(app) as c:
+        yield c
+
+def test_health_check(client):
+    """Test the root endpoint."""
+    response = client.get("/")
+    assert response.status_code == 200
+    assert response.json()["status"] == "ok"
+    assert response.json()["model_loaded"] is True
+
+def test_predict_spam(client):
+    """Test the prediction endpoint with a spam message."""
+    payload = {"text": "Congratulations! You've won a $1000 gift card. Click here."}
+    response = client.post("/predict", json=payload)
+    
+    assert response.status_code == 200
+    data = response.json()
+    assert data["prediction"] in ["spam", "ham"]
+
+def test_predict_ham(client):
+    """Test the prediction endpoint with a normal message."""
+    payload = {"text": "Hey, are we still meeting for lunch tomorrow?"}
+    response = client.post("/predict", json=payload)
+    
+    assert response.status_code == 200
+    data = response.json()
+    assert data["prediction"] == "ham" # Expecting 'ham' for this text