polyguard-rl/tests/test_provider_runtime.py

from app.common.enums import ActionType, DecisionMode, DoseBucket
from app.common.types import CandidateAction
from app.models.policy.provider_runtime import PolicyProviderRouter


def _candidate(idx: int, delta: float, uncertainty: float) -> CandidateAction:
    return CandidateAction(
        candidate_id=f"cand_{idx:02d}",
        mode=DecisionMode.REGIMEN_OPT,
        action_type=ActionType.KEEP_REGIMEN,
        target_drug=None,
        replacement_drug=None,
        dose_bucket=DoseBucket.NA,
        taper_days=None,
        monitoring_plan=None,
        estimated_safety_delta=delta,
        burden_delta=0.0,
        disease_stability_estimate=0.7,
        uncertainty_score=uncertainty,
        rationale_tags=["test"],
        required_monitoring=[],
        legality_precheck=True,
    )


def test_provider_router_returns_valid_candidate(monkeypatch) -> None:
    monkeypatch.setenv("POLYGUARD_ENABLE_ACTIVE_MODEL", "false")
    router = PolicyProviderRouter(ollama_model="qwen2.5:1.5b-instruct")
    candidates = [_candidate(1, 0.05, 0.6), _candidate(2, 0.2, 0.2)]
    selection = router.select_candidate(candidates, prompt={"task": "test"})
    assert selection.candidate_id in {"cand_01", "cand_02"}
    assert selection.provider
    assert router.model_status()["provider"] == "transformers"