tests/ml/test_model_manager.py

from pathlib import Path

import pytest

from app.core.exceptions import ServiceUnavailableError
from app.ml.model_manager import ModelManager


VALID_PAYLOAD = {
    "age": 45,
    "workclass": "Private",
    "fnlwgt": 250000,
    "education": "Masters",
    "education_num": 14,
    "marital_status": "Married-civ-spouse",
    "occupation": "Exec-managerial",
    "relationship": "Husband",
    "race": "White",
    "sex": "Male",
    "capital_gain": 15000,
    "capital_loss": 0,
    "hours_per_week": 50,
    "native_country": "United-States",
}


def test_model_manager_loads_real_artifact() -> None:
    manager = ModelManager(Path("app/ml/pipeline_produccion.pkl"))
    manager.load_model()

    assert manager.is_loaded is True
    assert manager.model_version


def test_model_manager_predict_one_real_artifact() -> None:
    manager = ModelManager(Path("app/ml/pipeline_produccion.pkl"))
    manager.load_model()

    prediction = manager.predict_one(VALID_PAYLOAD)

    assert prediction.label in {"<=50K", ">50K"}
    assert 0.0 <= prediction.probability <= 1.0
    assert len(prediction.raw_probabilities) == 2


def test_model_manager_rejects_missing_artifact(tmp_path) -> None:
    manager = ModelManager(tmp_path / "missing.pkl")
    with pytest.raises(ServiceUnavailableError):
        manager.load_model()