tests/test_api.py

# 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