Spaces:
Sleeping
Sleeping
File size: 1,793 Bytes
6172160 4904e85 6172160 | 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 | """Tests for the 2D dispatch visualizer."""
from __future__ import annotations
from pathlib import Path
from src.models import (
IncidentSeverity,
IncidentState,
IncidentStatus,
IncidentType,
State,
UnitState,
UnitStatus,
UnitType,
)
from src.visualizer.viewer import Viewer2D
def _state() -> State:
unit = UnitState(
unit_id="MED-1",
unit_type=UnitType.MEDIC,
status=UnitStatus.AVAILABLE,
location_x=10.0,
location_y=10.0,
assigned_incident_id=None,
eta_seconds=0.0,
crew_count=2,
)
inc = IncidentState(
incident_id="INC-001",
incident_type=IncidentType.CARDIAC_ARREST,
severity=IncidentSeverity.PRIORITY_1,
location_x=12.0,
location_y=14.0,
reported_at_step=0,
units_assigned=[],
status=IncidentStatus.PENDING,
survival_clock=100.0,
)
return State(
units={unit.unit_id: unit},
incidents={inc.incident_id: inc},
episode_id="ep",
step_count=0,
task_id="single_incident",
city_time=0.0,
metadata={"grid_size": [100, 100]},
)
def test_update_syncs_fields() -> None:
viewer = Viewer2D()
st = _state()
viewer.update(st)
assert viewer.step_count == 0
assert viewer.task_id == "single_incident"
assert "MED-1" in viewer.units
def test_render_png_bytes() -> None:
viewer = Viewer2D()
viewer.update(_state())
png = viewer.render()
assert png[:8] == b"\x89PNG\r\n\x1a\n"
def test_render_to_file(tmp_path: Path) -> None:
viewer = Viewer2D()
out = tmp_path / "frame.png"
viewer.render_to_file(str(out), _state())
assert out.exists()
assert out.read_bytes()[:8] == b"\x89PNG\r\n\x1a\n"
|