tests/test_env.py

"""Gymnasium API compliance tests for :class:`HospitalEnv`."""
from __future__ import annotations

import warnings

import numpy as np
import pytest
from gymnasium.utils.env_checker import check_env

from hospital_env import HospitalEnv


# ---------------------------------------------------------------------
# API compliance
# ---------------------------------------------------------------------
def test_env_passes_check_env() -> None:
    env = HospitalEnv()
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        check_env(env.unwrapped, skip_render_check=True)


def test_reset_returns_obs_info_tuple() -> None:
    env = HospitalEnv()
    out = env.reset(seed=0)
    assert isinstance(out, tuple) and len(out) == 2
    obs, info = out
    assert isinstance(obs, dict)
    assert isinstance(info, dict)
    # All observation components must match their declared space.
    for key, space in env.observation_space.spaces.items():
        assert key in obs
        assert obs[key].shape == space.shape
        assert obs[key].dtype == space.dtype


def test_step_returns_five_tuple() -> None:
    env = HospitalEnv()
    env.reset(seed=0)
    out = env.step(0)
    assert len(out) == 5
    obs, reward, terminated, truncated, info = out
    assert isinstance(obs, dict)
    assert isinstance(reward, float)
    assert isinstance(terminated, bool)
    assert isinstance(truncated, bool)
    assert isinstance(info, dict)


# ---------------------------------------------------------------------
# Reproducibility
# ---------------------------------------------------------------------
def test_seed_is_reproducible() -> None:
    env1 = HospitalEnv()
    env2 = HospitalEnv()
    obs1, _ = env1.reset(seed=42)
    obs2, _ = env2.reset(seed=42)
    for k in obs1:
        assert np.array_equal(obs1[k], obs2[k]), f"Mismatch in {k}"

    # Run identical action sequences and check identical rewards.
    rewards1, rewards2 = [], []
    for a in [0, 1, 11, 0, 2, 0, 3, 15, 0]:
        _, r1, *_ = env1.step(a)
        _, r2, *_ = env2.step(a)
        rewards1.append(r1)
        rewards2.append(r2)
    assert rewards1 == rewards2


# ---------------------------------------------------------------------
# Episode mechanics
# ---------------------------------------------------------------------
def test_full_random_episode_runs() -> None:
    env = HospitalEnv()
    env.reset(seed=0)
    steps = 0
    terminated = truncated = False
    while not (terminated or truncated):
        _, _, terminated, truncated, _ = env.step(env.action_space.sample())
        steps += 1
        assert steps <= HospitalEnv.MAX_STEPS + 1
    assert steps >= 1


def test_episode_truncates_at_max_steps() -> None:
    env = HospitalEnv(max_deaths=10_000)  # prevent early termination
    env.reset(seed=0)
    for _ in range(HospitalEnv.MAX_STEPS):
        _, _, terminated, truncated, _ = env.step(0)
        if terminated or truncated:
            break
    assert truncated
    assert env.current_step == HospitalEnv.MAX_STEPS


def test_invalid_action_yields_penalty() -> None:
    env = HospitalEnv()
    env.reset(seed=0)
    # Assigning to empty queue slots should be invalid and incur -1.
    _, reward, _, _, info = env.step(
        HospitalEnv.ACTION_ASSIGN_GENERAL_START + 9  # queue slot 9, unlikely populated
    )
    if info["invalid_action"]:
        assert reward <= 0.0


def test_action_space_size() -> None:
    env = HospitalEnv()
    assert env.action_space.n == HospitalEnv.NUM_ACTIONS == 39


# ---------------------------------------------------------------------
# Observation bounds
# ---------------------------------------------------------------------
def test_observations_within_space_bounds() -> None:
    env = HospitalEnv()
    obs, _ = env.reset(seed=7)
    rng = np.random.default_rng(7)
    for _ in range(50):
        action = int(rng.integers(0, env.action_space.n))
        obs, _, terminated, truncated, _ = env.step(action)
        assert env.observation_space.contains(obs), obs
        if terminated or truncated:
            break


def test_assign_queue_to_bed_actually_fills_bed() -> None:
    env = HospitalEnv(arrival_rate=3.0)  # guarantee arrivals
    env.reset(seed=1)
    # Walk until there's at least one patient in the queue
    for _ in range(5):
        env.step(0)
        if len(env.hospital.waiting_queue) > 0:
            break
    before = sum(b is not None for b in env.hospital.general_beds)
    env.step(HospitalEnv.ACTION_ASSIGN_GENERAL_START + 0)
    after = sum(b is not None for b in env.hospital.general_beds)
    assert after == before + 1


def test_action_mask_shape_and_noop_always_valid() -> None:
    env = HospitalEnv()
    env.reset(seed=0)
    mask = env.action_masks()
    assert mask.shape == (HospitalEnv.NUM_ACTIONS,)
    assert mask.dtype == bool
    # No-op (action 0) must always be valid.
    assert mask[HospitalEnv.ACTION_NOOP]


def test_action_mask_eliminates_invalid_actions() -> None:
    """A policy that only samples from the action mask must never trigger
    the invalid-action penalty for any of its actions."""
    env = HospitalEnv(arrival_rate=2.0)
    rng = np.random.default_rng(0)
    env.reset(seed=0)
    invalid = 0
    steps = 0
    for _ in range(400):
        mask = env.action_masks()
        valid_indices = np.where(mask)[0]
        action = int(rng.choice(valid_indices))
        _, _, term, trunc, info = env.step(action)
        steps += 1
        if info["invalid_action"]:
            invalid += 1
        if term or trunc:
            env.reset(seed=0)
    assert steps > 100
    assert invalid == 0, f"masked policy still produced {invalid} invalid actions"


def test_early_discharge_is_penalized_and_not_counted() -> None:
    """Regression test: closing the admit-then-discharge reward hack.

    Admitting a patient and immediately discharging them must:
      1. produce a strictly negative net reward, and
      2. NOT increment ``total_treated``.
    """
    env = HospitalEnv(arrival_rate=3.0)
    env.reset(seed=2)

    # Walk a couple steps so the queue has at least one patient.
    while len(env.hospital.waiting_queue) == 0:
        env.step(0)

    treated_before = env.hospital.total_treated
    _, r_admit, *_ = env.step(HospitalEnv.ACTION_ASSIGN_GENERAL_START)  # admit q[0]
    # Now bed[0] holds the freshly-admitted patient. Discharge it immediately.
    _, r_discharge, *_ = env.step(HospitalEnv.ACTION_DISCHARGE_GENERAL_START)
    treated_after = env.hospital.total_treated

    # Net reward of the admit + early-discharge must be negative.
    assert (r_admit + r_discharge) < 0, (r_admit, r_discharge)
    # Patient must NOT be counted as treated.
    assert treated_after == treated_before