Spaces:
Sleeping
Sleeping
File size: 10,941 Bytes
3e4c342 74965f9 3e4c342 74965f9 3e4c342 74965f9 3e4c342 74965f9 3e4c342 74965f9 3e4c342 74965f9 3e4c342 | 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 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 | """
Test API — proves step()/reset()/state() return correctly typed models.
"""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from fastapi.testclient import TestClient
from models import GridAction, GridObservation, ActionType
from server.nexusgrid_environment import NexusgridEnvironment
from server.app import app
class TestEnvironmentAPI:
"""Test the OpenEnv API interface."""
def _make_env(self) -> NexusgridEnvironment:
return NexusgridEnvironment()
def test_reset_returns_observation(self):
"""reset() should return a GridObservation."""
env = self._make_env()
obs = env.reset(seed=42, task_id=0)
assert isinstance(obs, GridObservation)
assert obs.tick == 0
assert obs.task_id == 0
assert obs.done is False
assert 58.0 <= obs.grid_frequency_hz <= 62.0
def test_reset_topology_present(self):
"""reset() observation should contain topology with nodes and edges."""
env = self._make_env()
obs = env.reset(seed=42, task_id=0)
topo = obs.topology_graph
assert "nodes" in topo
assert "edges" in topo
assert len(topo["nodes"]) == 20 # 20-node grid
assert len(topo["edges"]) == 40 # 40 edges
def test_step_returns_observation(self):
"""step() should return a GridObservation with reward and done."""
env = self._make_env()
env.reset(seed=42, task_id=0)
action = GridAction(
action_type=ActionType.DISPATCH_GENERATION,
node_id="NODE_01",
mw=100.0,
)
obs = env.step(action)
assert isinstance(obs, GridObservation)
assert obs.tick == 1
assert obs.grid_frequency_hz > 0
assert "reward_breakdown" in obs.metadata
assert "rubric_breakdown" in obs.metadata
def test_step_advance_tick(self):
"""advance_tick should increment the tick counter."""
env = self._make_env()
env.reset(seed=42, task_id=1)
action = GridAction(action_type=ActionType.ADVANCE_TICK)
obs = env.step(action)
assert obs.tick == 1
def test_state_returns_state(self):
"""state property should return an OpenEnv State."""
env = self._make_env()
env.reset(seed=42, task_id=0)
state = env.state
assert state.episode_id is not None
assert state.step_count == 0
def test_health_endpoint_returns_200(self):
"""The standalone health endpoint should respond immediately."""
client = TestClient(app)
response = client.get("/health")
assert response.status_code == 200
assert response.json()["status"] == "healthy"
def test_web_dashboard_route_available(self):
"""The mounted dashboard route should remain available for HF Spaces."""
client = TestClient(app)
response = client.get("/web", follow_redirects=True)
assert response.status_code == 200
def test_websocket_reset_step_state_flow(self):
"""The persistent OpenEnv WebSocket API should support reset, step, and state."""
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
ws.send_json({"type": "reset", "data": {"seed": 42, "task_id": 0}})
reset_message = ws.receive_json()
assert reset_message["type"] == "observation"
assert reset_message["data"]["observation"]["tick"] == 0
ws.send_json(
{
"type": "step",
"data": {
"action_type": "dispatch_generation",
"node_id": "NODE_01",
"mw": 100,
},
}
)
step_message = ws.receive_json()
assert step_message["type"] == "observation"
assert step_message["data"]["observation"]["tick"] == 1
ws.send_json({"type": "state"})
state_message = ws.receive_json()
assert state_message["type"] == "state"
assert state_message["data"]["step_count"] == 1
def test_done_when_frequency_below_59(self):
"""Episode should terminate when frequency drops below 59.0Hz."""
env = self._make_env()
env.reset(seed=42, task_id=5) # Black start — frequency starts at 59.0
# In black start, frequency can drop. Let's force it by doing nothing.
# The initial frequency is 59.0 for black start
action = GridAction(action_type=ActionType.ADVANCE_TICK)
done = False
for _ in range(5):
obs = env.step(action)
if obs.done:
done = True
break
# Black start starts at 59.0 which is already at termination boundary
# The env should detect this
assert done or obs.grid_frequency_hz >= 59.0
def test_smoke_test_scores_1(self):
"""Task 0 with valid dispatch should score 1.0."""
env = self._make_env()
env.reset(seed=42, task_id=0)
action = GridAction(
action_type=ActionType.DISPATCH_GENERATION,
node_id="NODE_01",
mw=100.0,
)
env.step(action)
score = env.get_score()
assert score == 1.0, f"Smoke test should score 1.0, got {score}"
def test_dispatch_updates_frequency_without_advance_tick(self):
"""Control actions should update the observable grid frequency immediately."""
env = self._make_env()
obs = env.reset(seed=42, task_id=1)
initial_frequency = obs.grid_frequency_hz
obs = env.step(
GridAction(
action_type=ActionType.DISPATCH_GENERATION,
node_id="NODE_04",
mw=200.0,
)
)
assert obs.grid_frequency_hz != initial_frequency
def test_cascade_overload_reset_exposes_overloaded_lines(self):
"""Task 2 should surface the reroute overload in the initial observation."""
env = self._make_env()
obs = env.reset(seed=42, task_id=2)
overloaded = [
edge["id"]
for edge in obs.topology_graph["edges"]
if edge["status"] == "LIVE" and edge["current_load_mw"] >= 0.95 * edge["capacity_mw"]
]
assert "LINE_29" in overloaded
def test_cascade_overload_persists_without_immediate_action(self):
"""Task 2 should remain overloaded after a no-op tick until the agent isolates it."""
env = self._make_env()
env.reset(seed=42, task_id=2)
obs = env.step(GridAction(action_type=ActionType.ADVANCE_TICK))
overloaded = [
edge["id"]
for edge in obs.topology_graph["edges"]
if edge["status"] == "LIVE" and edge["current_load_mw"] >= 0.95 * edge["capacity_mw"]
]
assert "LINE_29" in overloaded
def test_black_start_can_reenergize_critical_loads(self):
"""Closing breakers from the hydro island should energize downstream critical loads."""
env = self._make_env()
env.reset(seed=42, task_id=5)
actions = [
GridAction(action_type=ActionType.DISPATCH_GENERATION, node_id="NODE_01", mw=600.0),
GridAction(action_type=ActionType.TOGGLE_CIRCUIT_BREAKER, edge_id="LINE_02", status="CLOSED"),
GridAction(action_type=ActionType.TOGGLE_CIRCUIT_BREAKER, edge_id="LINE_28", status="CLOSED"),
GridAction(action_type=ActionType.DISPATCH_GENERATION, node_id="NODE_17", mw=400.0),
GridAction(action_type=ActionType.TOGGLE_CIRCUIT_BREAKER, edge_id="LINE_22", status="CLOSED"),
]
obs = None
for action in actions:
obs = env.step(action)
assert obs is not None
assert env._engine.nodes["NODE_03"]["energized"] is True
assert env._engine.nodes["NODE_18"]["energized"] is True
assert env._engine.nodes["NODE_18"]["consumption_mw"] > 0.0
assert env.get_score() > 0.0
def test_full_restoration_tick_records_first_recovery(self):
"""Task 2 should keep the first tick that crosses the restoration threshold."""
env = self._make_env()
env.reset(seed=42, task_id=2)
env.step(GridAction(action_type=ActionType.TOGGLE_CIRCUIT_BREAKER, edge_id="LINE_29", status="OPEN"))
assert env._get_full_restoration_tick() is None
env.step(GridAction(action_type=ActionType.DISPATCH_GENERATION, node_id="NODE_09", mw=100.0))
assert env._get_full_restoration_tick() is None
env.step(GridAction(action_type=ActionType.DISPATCH_GENERATION, node_id="NODE_13", mw=100.0))
assert env._get_full_restoration_tick() == 3
def test_reset_idempotent(self):
"""Calling reset(42) multiple times should produce identical observations."""
env = self._make_env()
obs1 = env.reset(seed=42, task_id=0)
obs2 = env.reset(seed=42, task_id=0)
assert obs1.grid_frequency_hz == obs2.grid_frequency_hz
assert obs1.tick == obs2.tick
assert obs1.task_id == obs2.task_id
def test_done_observation_includes_episode_summary(self):
"""Final observations should expose a compact episode summary for training logs."""
env = self._make_env()
env.reset(seed=42, task_id=0)
obs = env.step(
GridAction(
action_type=ActionType.DISPATCH_GENERATION,
node_id="NODE_01",
mw=100.0,
)
)
if not obs.done:
obs = env.step(GridAction(action_type=ActionType.ADVANCE_TICK))
if not obs.done:
obs = env.step(GridAction(action_type=ActionType.ADVANCE_TICK))
assert obs.done is True
assert "episode_summary" in obs.metadata
def test_all_action_types_accepted(self):
"""All action types should be accepted without crashing."""
env = self._make_env()
env.reset(seed=42, task_id=1)
actions = [
GridAction(action_type=ActionType.ADVANCE_TICK),
GridAction(action_type=ActionType.DISPATCH_GENERATION, node_id="NODE_01", mw=50),
GridAction(action_type=ActionType.TOGGLE_CIRCUIT_BREAKER, edge_id="LINE_01", status="OPEN"),
GridAction(action_type=ActionType.RUN_STATE_ESTIMATION, subgraph=["NODE_01", "NODE_02"]),
GridAction(action_type=ActionType.QUARANTINE_SCADA_NODE, node_id="NODE_01"),
]
for action in actions:
obs = env.step(action)
assert isinstance(obs, GridObservation)
def test_weather_present(self):
"""Observation should include weather data."""
env = self._make_env()
obs = env.reset(seed=42, task_id=0)
assert len(obs.weather_forecast_matrix) > 0
assert obs.weather_summary != ""
|