tests/test_openenv.py

"""Tests for the OpenEnv adapter layer (HospitalOpenEnv + app:app).

These exercise the OpenEnv contract directly without spinning up a
real HTTP server, so they're fast and deterministic.
"""
from __future__ import annotations

import pytest

# Skip the whole module if openenv-core isn't installed (minimal install).
pytest.importorskip("openenv.core")
pytest.importorskip("fastapi")

from fastapi.testclient import TestClient  # noqa: E402

from hospital_env import (  # noqa: E402
    HospitalAction,
    HospitalObservation,
    HospitalOpenEnv,
    HospitalState,
)


# ---------------------------------------------------------------------
# HospitalOpenEnv unit tests
# ---------------------------------------------------------------------
def test_openenv_reset_returns_observation() -> None:
    env = HospitalOpenEnv()
    obs = env.reset(seed=0)
    assert isinstance(obs, HospitalObservation)
    assert obs.done is False
    assert obs.reward is None
    assert len(obs.bed_occupancy) == 20
    assert len(obs.icu_occupancy) == 5
    assert len(obs.ventilator_status) == 3
    assert len(obs.action_mask) == 39
    assert obs.time_step == 0


def test_openenv_step_advances_time() -> None:
    env = HospitalOpenEnv()
    env.reset(seed=0)
    obs = env.step(HospitalAction(action=0))
    assert isinstance(obs, HospitalObservation)
    assert obs.time_step == 1
    assert isinstance(obs.reward, float)


def test_openenv_state_property() -> None:
    env = HospitalOpenEnv()
    env.reset(seed=0, episode_id="test-ep")
    for _ in range(10):
        env.step(HospitalAction(action=0))
    state = env.state
    assert isinstance(state, HospitalState)
    assert state.episode_id == "test-ep"
    assert state.step_count == 10
    assert state.queue_len >= 0
    assert state.scenario is None  # default config


def test_openenv_metadata() -> None:
    env = HospitalOpenEnv()
    md = env.get_metadata()
    assert "Hospital" in md.name
    assert md.version is not None
    assert len(md.description) > 30


def test_openenv_scenario_dispatch() -> None:
    env = HospitalOpenEnv(scenario="surge")
    obs = env.reset(seed=0)
    assert obs.time_step == 0
    state = env.state
    assert state.scenario == "surge"


def test_openenv_unknown_scenario_raises() -> None:
    with pytest.raises(ValueError, match="Unknown scenario"):
        HospitalOpenEnv(scenario="not_a_real_scenario")


def test_openenv_close_is_idempotent() -> None:
    env = HospitalOpenEnv()
    env.reset(seed=0)
    env.close()
    env.close()  # second call must not raise


# ---------------------------------------------------------------------
# FastAPI app smoke tests (in-process via TestClient)
# ---------------------------------------------------------------------
@pytest.fixture(scope="module")
def app_client() -> TestClient:
    from app import app

    return TestClient(app)


def test_app_health(app_client: TestClient) -> None:
    r = app_client.get("/health")
    assert r.status_code == 200
    assert r.json()["status"] == "healthy"


def test_app_metadata(app_client: TestClient) -> None:
    r = app_client.get("/metadata")
    assert r.status_code == 200
    md = r.json()
    assert "Hospital" in md["name"]
    assert md["version"] == "1.0.0"


def test_app_schema_has_action_observation_state(app_client: TestClient) -> None:
    r = app_client.get("/schema")
    assert r.status_code == 200
    s = r.json()
    assert "action" in s
    assert "observation" in s
    assert "state" in s
    obs_props = s["observation"].get("properties", {})
    for key in (
        "bed_occupancy",
        "icu_occupancy",
        "ventilator_status",
        "waiting_queue",
        "time_step",
        "stats",
        "action_mask",
    ):
        assert key in obs_props, f"missing observation field: {key}"


def test_app_reset_via_http(app_client: TestClient) -> None:
    r = app_client.post("/reset", json={"seed": 0})
    assert r.status_code == 200
    body = r.json()
    obs = body.get("observation", body)
    assert obs["time_step"] == 0
    assert len(obs["action_mask"]) == 39


def test_app_step_via_http(app_client: TestClient) -> None:
    r = app_client.post("/step", json={"action": {"action": 0}})
    assert r.status_code == 200
    body = r.json()
    obs = body.get("observation", body)
    # /step is stateless on REST: each call constructs a fresh env, so we
    # only assert the protocol shape, not stateful progression.
    assert "time_step" in obs
    assert "action_mask" in obs


def test_app_step_rejects_invalid_action_value(app_client: TestClient) -> None:
    r = app_client.post("/step", json={"action": {"action": 999}})
    # Pydantic validation should reject (action <= 38 constraint).
    assert r.status_code in (400, 422)