""" test_weather_forecast_env_outcome_codes.py =========================================== Spec suite for the LOCKED episode-lifecycle outcome-code contract on WeatherForecastEnv. Written against the contract, not against an implementation that exists yet: `info["outcome_code"]` is not populated by the current step()/reset() code, so every RESET_OK / INSPECT_OK / TERMINATED_* / PENALTY_* test in this file is EXPECTED TO FAIL until the outcome_code patch lands. That failure is the point (build-order rule: contract -> tests -> code). The four ERROR_* tests assert on raised exceptions, which already exist in the current code and should pass today. Locked contract this file tests against ---------------------------------------- Returned paths (info["outcome_code"]): RESET_OK INSPECT_OK -- valid inspect, episode continues INSPECT_AND_BUDGET_EXHAUSTED -- valid inspect that also hits max_steps TERMINATED_EARLY -- explicit terminate action TERMINATED_BUDGET -- budget exhausted via terminate action OR via a penalty branch (accepted documented asymmetry vs the inspect case above) PENALTY_REVISIT -- episode not yet terminated PENALTY_PADDING -- episode not yet terminated PENALTY_INVALID -- episode not yet terminated Raised, never coded (real exceptions, not outcome codes): step() before reset() -> RuntimeError injected context length / missing zone_obs lists -> ValueError n_active > max_zones -> ValueError NaNSafetyWrapper nan_limit exceeded -> RuntimeError UNDEFINED, intentionally not tested here (per the contract): behaviour of a second step() after any terminal outcome_code whether a future wrapper ever converts raises into returned codes FIXTURES -------- Grounded directly in zone_observation.py (read in full, not inferred): ForecastConfig, EpisodeContext, ZoneObs, ForecastResult, and the make_synthetic_* helpers. Fixture-construction bugs found and fixed while cross-checking against the real source (see conversation notes): - EpisodeContext requires real obs/forecast objects; obs=None fails in __post_init__ with AttributeError, not the ValueError being tested for. - Injecting a single-zone EpisodeContext (zone_ids length 1, zone_obs/ zone_forecasts provided) into a 2-slot env is what actually reaches the padding branch -- previously skipped as "not constructible". - A legacy EpisodeContext with zone_ids length 2 but empty zone_obs/ zone_forecasts lists reaches WeatherForecastEnv's "zone payload length mismatch" ValueError via resolved_zone_obs() falling back to [obs] (length 1 != n_active). - Corrupting belief_map[action] before stepping into that same action gets silently healed by _compute_zone_refinement_reward (NaN comparisons are always False, so max(nan - rate, signal) == signal). Must corrupt a zone other than the one being inspected. """ from __future__ import annotations import numpy as np import pytest from zone_observation import ( EpisodeContext, ForecastConfig, make_synthetic_forecast_result, make_synthetic_zone_obs, ) from weather_forecast_env import ( WeatherForecastEnv, NaNSafetyWrapper, make_weather_env, ) # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- def _cfg(**overrides) -> ForecastConfig: base = dict( n_zones=2, horizon_days=3, max_steps=1, prior_belief=0.12, clean_episode_ratio=0.0, # force a hazard every episode -> deterministic tests event_spatial_correlation=0.85, shuffle_zone_order=False, # deterministic action->zone mapping for these tests seed=7, ) base.update(overrides) return ForecastConfig(**base) def _zone_obs_and_forecast(zone_id: str, cfg: ForecastConfig, seed: int, **event_flags): obs = make_synthetic_zone_obs(zone_id, seed=seed, **event_flags) fc = make_synthetic_forecast_result( zone_id, valid_time=obs.valid_time, horizon_days=cfg.horizon_days, seed=seed, **{k: v for k, v in event_flags.items() if k in ("drought", "flood")}, ) return obs, fc @pytest.fixture def env_triage(): return WeatherForecastEnv(_cfg(n_zones=2, max_steps=1)) @pytest.fixture def env_full_tour(): return WeatherForecastEnv(_cfg(n_zones=2, max_steps=5)) def _outcome(info: dict) -> str: assert "outcome_code" in info, ( "info['outcome_code'] is not populated by the current code. " "This is expected until the outcome-code patch is implemented; " "this suite is the spec for that patch, not a report of current " "passing behaviour." ) return info["outcome_code"] # --------------------------------------------------------------------------- # RESET_OK # --------------------------------------------------------------------------- def test_reset_ok(env_full_tour): obs, info = env_full_tour.reset(seed=1) assert _outcome(info) == "RESET_OK" def test_reset_ok_is_deterministic_for_same_seed_and_config(): e1 = WeatherForecastEnv(_cfg(seed=42)) e2 = WeatherForecastEnv(_cfg(seed=42)) obs1, info1 = e1.reset(seed=42) obs2, info2 = e2.reset(seed=42) np.testing.assert_array_equal(obs1["zone_belief"], obs2["zone_belief"]) np.testing.assert_array_equal(obs1["forecast_precip"], obs2["forecast_precip"]) assert info1["zone_ids"] == info2["zone_ids"] # --------------------------------------------------------------------------- # INSPECT_OK vs INSPECT_AND_BUDGET_EXHAUSTED (the case Option 2 exists for) # --------------------------------------------------------------------------- def test_inspect_ok_when_episode_continues(env_full_tour): env_full_tour.reset(seed=1) obs, reward, terminated, truncated, info = env_full_tour.step(0) assert terminated is False assert truncated is False assert _outcome(info) == "INSPECT_OK" def test_inspect_and_budget_exhausted_under_triage(env_triage): env_triage.reset(seed=1) obs, reward, terminated, truncated, info = env_triage.step(0) assert terminated is True assert truncated is False assert _outcome(info) == "INSPECT_AND_BUDGET_EXHAUSTED" assert info["zone_id"] is not None assert info["visited"][0] == True # noqa: E712 (explicit bool array check) def test_inspect_ok_zero_inspect_null_control(env_full_tour): obs0, _ = env_full_tour.reset(seed=1) belief_before = obs0["zone_belief"].copy() obs1, reward, terminated, truncated, info = env_full_tour.step( env_full_tour.terminate_action ) np.testing.assert_array_equal(obs1["zone_belief"], belief_before) assert _outcome(info) == "TERMINATED_EARLY" # --------------------------------------------------------------------------- # TERMINATED_EARLY / TERMINATED_BUDGET # --------------------------------------------------------------------------- def test_terminated_early_on_explicit_terminate(env_full_tour): env_full_tour.reset(seed=3) _, _, terminated, _, info = env_full_tour.step(env_full_tour.terminate_action) assert terminated is True assert _outcome(info) == "TERMINATED_EARLY" def test_terminated_budget_via_penalty_branch(env_triage): env_triage.reset(seed=5) invalid_action = env_triage.max_zones + 5 # out of legal range _, _, terminated, _, info = env_triage.step(invalid_action) assert terminated is True assert _outcome(info) == "TERMINATED_BUDGET" assert info.get("invalid_action") is True def test_terminate_and_budget_exhaust_agree_on_reward_shape(): e1 = WeatherForecastEnv(_cfg(n_zones=2, max_steps=5, seed=9)) e1.reset(seed=9) e1.step(0) _, r1, term1, _, info1 = e1.step(e1.terminate_action) e2 = WeatherForecastEnv(_cfg(n_zones=2, max_steps=2, seed=9)) e2.reset(seed=9) e2.step(0) _, r2, term2, _, info2 = e2.step(1) assert term1 is True and term2 is True assert _outcome(info1) == "TERMINATED_EARLY" assert _outcome(info2) == "INSPECT_AND_BUDGET_EXHAUSTED" assert "believed_p" in info1 and "believed_p" in info2 # --------------------------------------------------------------------------- # PENALTY_REVISIT / PENALTY_PADDING / PENALTY_INVALID (episode continues) # --------------------------------------------------------------------------- def test_penalty_revisit_when_not_terminated(env_full_tour): env_full_tour.reset(seed=11) env_full_tour.step(0) # first visit, legal _, reward, terminated, _, info = env_full_tour.step(0) # revisit assert terminated is False assert _outcome(info) == "PENALTY_REVISIT" assert reward < 0 def test_penalty_padding_when_not_terminated(): cfg = _cfg(n_zones=2, max_steps=5) env = WeatherForecastEnv(cfg) obs0, fc0 = _zone_obs_and_forecast("z0", cfg, seed=31, drought=True) single_zone_context = EpisodeContext( obs=obs0, forecast=fc0, config=cfg, zone_ids=["z0"], zone_obs=[obs0], zone_forecasts=[fc0], ) env.reset(seed=1, options={"context": single_zone_context}) _, reward, terminated, _, info = env.step(1) # slot 1: padding, n_active=1 assert terminated is False assert _outcome(info) == "PENALTY_PADDING" assert reward < 0 def test_penalty_invalid_when_not_terminated(env_full_tour): env_full_tour.reset(seed=17) invalid_action = env_full_tour.max_zones + 5 _, reward, terminated, _, info = env_full_tour.step(invalid_action) assert terminated is False assert _outcome(info) == "PENALTY_INVALID" assert reward < 0 # --------------------------------------------------------------------------- # Shuffle mapping (test scenario #7) # --------------------------------------------------------------------------- def test_shuffle_changes_action_to_zone_mapping_not_belief_semantics(): cfg_shuffled = _cfg(n_zones=2, max_steps=5, shuffle_zone_order=True, seed=21) env = WeatherForecastEnv(cfg_shuffled) obs, info = env.reset(seed=21) zone_ids_at_reset = list(info["zone_ids"]) _, _, _, _, info_after_step = env.step(0) assert info_after_step["zone_ids"] == zone_ids_at_reset # --------------------------------------------------------------------------- # Raised, never coded (real exceptions -- should pass against current code) # --------------------------------------------------------------------------- def test_step_before_reset_raises_runtime_error(): env = WeatherForecastEnv(_cfg()) with pytest.raises(RuntimeError): env.step(0) def test_zone_overflow_raises_value_error(): cfg = _cfg(n_zones=2, max_steps=5) env = WeatherForecastEnv(cfg) obs0, fc0 = _zone_obs_and_forecast("z0", cfg, seed=33) bad_context = EpisodeContext( obs=obs0, forecast=fc0, config=cfg, zone_ids=["z0", "z1", "z2"], ) with pytest.raises(ValueError): env.reset(seed=1, options={"context": bad_context}) def test_injection_length_mismatch_raises_value_error(): cfg = _cfg(n_zones=2, max_steps=5) env = WeatherForecastEnv(cfg) obs0, fc0 = _zone_obs_and_forecast("z0", cfg, seed=35) legacy_multi_zone_context = EpisodeContext( obs=obs0, forecast=fc0, config=cfg, zone_ids=["z0", "z1"], ) with pytest.raises(ValueError): env.reset(seed=1, options={"context": legacy_multi_zone_context}) def test_nan_limit_exceeded_raises_runtime_error(): cfg = _cfg(n_zones=2, max_steps=5) inner = WeatherForecastEnv(cfg) wrapped = NaNSafetyWrapper(inner, nan_limit=0) wrapped.reset(seed=1) inner._belief_map[1] = float("nan") with pytest.raises(RuntimeError): wrapped.step(0) # --------------------------------------------------------------------------- # max_steps=1 (triage) vs full tour -- test scenario #8 # --------------------------------------------------------------------------- def test_triage_vs_full_tour_are_not_interchangeable(env_triage, env_full_tour): assert env_triage.config.max_steps < env_triage.max_zones assert env_full_tour.config.max_steps >= env_full_tour.max_zones