tests/test_env_step.py

from app.env.env_core import PolyGuardEnv
import pytest


def test_env_step_returns_valid_reward() -> None:
    env = PolyGuardEnv()
    env.reset(seed=9, difficulty="easy")
    action = env.get_legal_actions()[0]
    _, reward, done, info = env.step(action)
    assert 0.001 <= reward <= 0.999
    assert isinstance(done, bool)
    assert "safety_report" in info


def test_env_step_rejects_unknown_action_keys() -> None:
    env = PolyGuardEnv()
    env.reset(seed=11, difficulty="easy")
    action = env.get_legal_actions()[0]
    action["unexpected_key"] = "not_allowed"
    with pytest.raises(Exception):
        env.step(action)