polyguard-rl/tests/test_graph_infer.py

from __future__ import annotations

import pickle
from pathlib import Path

from app.models.graph.infer import infer_graph_risk


class BrokenProbabilityModel:
    def predict_proba(self, _encoded):
        raise AttributeError("multi_class")


class FakeLabelBinarizer:
    classes_ = ["dizziness"]


def test_graph_inference_uses_deterministic_fallback_for_stale_model(tmp_path: Path) -> None:
    model_path = tmp_path / "graph_model.pkl"
    with model_path.open("wb") as handle:
        pickle.dump(
            {
                "severe_model": BrokenProbabilityModel(),
                "side_model": BrokenProbabilityModel(),
                "mlb": FakeLabelBinarizer(),
            },
            handle,
        )

    risk = infer_graph_risk(["warfarin", "aspirin"], model_path=model_path)

    assert 0.0 <= risk["severe_alert_probability"] <= 1.0
    assert isinstance(risk["side_effect_probs"], dict)
    assert "warfarin__aspirin" in risk["pairwise_ddi_severity"]