test_injection_guard.py

#!/usr/bin/env python3
"""
test_injection_guard.py
=======================
Prove injected real-eval path:
  - n_zones=1 succeeds (legacy single obs/forecast)
  - n_zones>1 without zone_obs lists raises (no silent padding)
  - n_zones>1 WITH explicit zone_obs/zone_forecasts lists succeeds (Blocker B)
  - synthetic training reset still works at n_zones=3
"""
from __future__ import annotations

import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))


def main() -> int:
    from zone_observation import (
        EpisodeContext,
        ForecastConfig,
        make_synthetic_zone_obs,
        make_synthetic_forecast_result,
    )
    from weather_forecast_env import make_weather_env

    print("test_injection_guard")

    # n_zones=1 + inject must succeed
    cfg1 = ForecastConfig(n_zones=1, max_steps=4, horizon_days=30)
    env1 = make_weather_env(cfg1, use_nan_wrapper=True)
    obs = make_synthetic_zone_obs("karawang_rice", drought=True, seed=1)
    fc = make_synthetic_forecast_result(
        zone_id="karawang_rice", valid_time=obs.valid_time, drought=True, seed=1
    )
    ctx1 = EpisodeContext(
        zone_ids=["karawang_rice"],
        obs=obs,
        forecast=fc,
        config=cfg1,
    )
    o, info = env1.reset(options={"context": ctx1})
    assert o["zone_belief"].shape[0] >= 1
    print("  n_zones=1 inject OK")

    # n_zones=3 + inject WITHOUT zone_obs lists must raise (not pad)
    cfg3 = ForecastConfig(n_zones=3, max_steps=6, horizon_days=30)
    env3 = make_weather_env(cfg3, use_nan_wrapper=True)
    ctx_bad = EpisodeContext(
        zone_ids=["karawang_rice", "indramayu_rice", "central_java_rice"],
        obs=obs,
        forecast=fc,
        config=cfg3,
    )
    raised = False
    try:
        env3.reset(options={"context": ctx_bad})
    except ValueError as e:
        raised = True
        msg = str(e)
        assert (
            "n_zones=1 only" in msg
            or "without explicit zone_obs" in msg
            or "parallel to zone_ids" in msg
            or "refuses n_zones>1" in msg
        ), msg
        print(f"  n_zones=3 inject without lists raised as expected: {msg[:90]}...")
    if not raised:
        print("  FAIL: n_zones=3 inject did not raise — padding bug may still be live")
        return 1

    # n_zones=3 + inject WITH explicit per-zone lists must succeed (Blocker B)
    zids = ["karawang_rice", "indramayu_rice", "central_java_rice"]
    z_obs = [
        make_synthetic_zone_obs(z, drought=(i == 0), flood=(i == 1), seed=10 + i)
        for i, z in enumerate(zids)
    ]
    # Align valid_time so forecasts share a coherent episode clock
    vt = z_obs[0].valid_time
    for zo in z_obs:
        zo.valid_time = vt
    z_fc = [
        make_synthetic_forecast_result(
            zone_id=z, valid_time=vt, drought=(i == 0), flood=(i == 1), seed=20 + i
        )
        for i, z in enumerate(zids)
    ]
    ctx_ok = EpisodeContext(
        zone_ids=zids,
        obs=z_obs[0],
        forecast=z_fc[0],
        config=cfg3,
        zone_obs=z_obs,
        zone_forecasts=z_fc,
    )
    o_ok, info_ok = env3.reset(options={"context": ctx_ok})
    assert o_ok["zone_belief"].shape[0] >= 3
    # Beliefs should differ across zones (different event flags / noise)
    zb = o_ok["zone_belief"][:3]
    print(f"  n_zones=3 inject with lists OK  zone_belief={zb.tolist()}")

    # Training path (no inject) must still work at n_zones=3
    o3, _ = env3.reset(seed=0)
    assert o3["zone_belief"].shape[0] >= 3
    print("  n_zones=3 synthetic training reset OK")
    print("All injection-guard tests passed.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())