Spaces:
Sleeping
Sleeping
File size: 3,268 Bytes
9ea8001 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool
from app.main import app
from app.core.database import Base, get_db
from app.core.config import settings
# Setup in-memory SQLite database for testing
SQLALCHEMY_DATABASE_URL = "sqlite:///:memory:"
engine = create_engine(
SQLALCHEMY_DATABASE_URL,
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def override_get_db():
try:
db = TestingSessionLocal()
yield db
finally:
db.close()
app.dependency_overrides[get_db] = override_get_db
client = TestClient(app)
@pytest.fixture(scope="module")
def test_db():
# Create tables
Base.metadata.create_all(bind=engine)
yield
# Drop tables
Base.metadata.drop_all(bind=engine)
def test_health_check():
response = client.get("/health")
assert response.status_code == 200
assert response.json() == {"status": "healthy"}
def test_predict_unauthorized():
response = client.post("/predict", json={})
assert response.status_code == 422
def test_predict_invalid_key():
response = client.post("/predict", headers={"X-API-KEY": "wrong_key"}, json={})
assert response.status_code == 401
def test_predict_success(test_db):
# Valid input data based on updated schema
payload = {
"age": 30,
"genre": "M",
"revenu_mensuel": 5000,
"statut_marital": "Célibataire",
"departement": "R&D",
"poste": "Ingénieur",
"nombre_experiences_precedentes": 2,
"nombre_heures_travailless": 40,
"annee_experience_totale": 5,
"annees_dans_l_entreprise": 2,
"annees_dans_le_poste_actuel": 1,
"satisfaction_employee_environnement": 3,
"note_evaluation_precedente": 3,
"niveau_hierarchique_poste": 2,
"satisfaction_employee_nature_travail": 3,
"satisfaction_employee_equipe": 4,
"satisfaction_employee_equilibre_pro_perso": 3,
"note_evaluation_actuelle": 3,
"heure_supplementaires": "Non",
"augementation_salaire_precedente": "10-15%",
"nombre_participation_pee": 0,
"nb_formations_suivies": 1,
"nombre_employee_sous_responsabilite": 0,
"distance_domicile_travail": 10,
"niveau_education": 3,
"domaine_etude": "Sciences",
"ayant_enfants": "Non",
"frequence_deplacement": "Rare",
"annees_depuis_la_derniere_promotion": 1,
"annes_sous_responsable_actuel": 1
}
response = client.post(
"/predict",
headers={"X-API-KEY": settings.API_KEY},
json=payload
)
if response.status_code != 200:
print(f"Response error: {response.json()}")
assert response.status_code == 200
data = response.json()
assert "prediction" in data
assert "probability" in data
# Verify DB insertion
db = TestingSessionLocal()
log = db.query(Base.metadata.tables["prediction_logs"]).first()
assert log is not None
assert log.age == 30
assert log.genre == "M"
db.close()
|