File size: 3,711 Bytes
976eb45 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | #!/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())
|