polyguard-rl/tests/test_future_subenvs.py

from app.common.enums import ActionType
from app.env.env_core import PolyGuardEnv


def _candidate_types(env: PolyGuardEnv) -> set[str]:
    return {item["action_type"] for item in env.get_candidate_actions()}


def test_web_search_subenv_exposes_fetch_evidence_action() -> None:
    env = PolyGuardEnv()
    env.reset(seed=201, difficulty="hard", sub_environment="WEB_SEARCH_MISSING_DATA")
    assert ActionType.FETCH_EXTERNAL_EVIDENCE.value in _candidate_types(env)


def test_alternative_subenv_exposes_alternative_action() -> None:
    env = PolyGuardEnv()
    env.reset(seed=202, difficulty="medium", sub_environment="ALTERNATIVE_SUGGESTION")
    assert ActionType.RECOMMEND_ALTERNATIVE.value in _candidate_types(env)


def test_new_drug_subenv_exposes_component_decomposition_action() -> None:
    env = PolyGuardEnv()
    env.reset(seed=203, difficulty="hard", sub_environment="NEW_DRUG_DECOMPOSITION")
    assert ActionType.DECOMPOSE_NEW_DRUG.value in _candidate_types(env)