refactor(simulator): extract habitat simulator and visualization modules

mini-nav/tests/test_room_views.py

@@ -0,0 +1,108 @@
+from types import SimpleNamespace
+
+import pytest
+
+from simulator import collect_room_views_by_room
+
+
+class _FakeAgent:
+    def __init__(self):
+        self.positions = []
+
+    def set_state(self, state):
+        self.positions.append(state.position)
+
+
+class _FakeSimulator:
+    def __init__(self):
+        self._frame_index = 0
+        self.actions = []
+
+    def get_sensor_observations(self):
+        observations = {
+            "color_sensor": f"frame_{self._frame_index}",
+            "depth_sensor": f"depth_{self._frame_index}",
+        }
+        self._frame_index += 1
+        return observations
+
+    def step(self, action_name):
+        self.actions.append(action_name)
+
+
+class _FakeAgentState:
+    def __init__(self):
+        self.position = None
+
+
+def test_collect_room_views_by_room_collects_grouped_frames_with_single_outer_progress():
+    track_calls = []
+
+    def fake_track(iterable, description):
+        track_calls.append(description)
+        return iterable
+
+    agent = _FakeAgent()
+    sim = _FakeSimulator()
+    room_nodes = [
+        SimpleNamespace(room_id="room_00", center=[1.0, 2.0, 3.0]),
+        SimpleNamespace(room_id="room_01", center=[4.0, 5.0, 6.0]),
+    ]
+    fake_habitat = SimpleNamespace(AgentState=_FakeAgentState)
+
+    room_views = collect_room_views_by_room(
+        agent=agent,
+        sim=sim,
+        room_nodes=room_nodes,
+        views_per_room=3,
+        habitat_sim_module=fake_habitat,
+        progress_track=fake_track,
+    )
+
+    assert track_calls == ["Collecting room views"]
+    assert room_views == {
+        "room_00": ["frame_0", "frame_1", "frame_2"],
+        "room_01": ["frame_3", "frame_4", "frame_5"],
+    }
+    assert sim.actions == [
+        "turn_left",
+        "turn_left",
+        "turn_left",
+        "turn_left",
+        "turn_left",
+        "turn_left",
+    ]
+    assert agent.positions == [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+
+
+def test_collect_room_views_by_room_uses_custom_sensor_and_turn_action():
+    agent = _FakeAgent()
+    sim = _FakeSimulator()
+    room_nodes = [SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])]
+    fake_habitat = SimpleNamespace(AgentState=_FakeAgentState)
+
+    room_views = collect_room_views_by_room(
+        agent=agent,
+        sim=sim,
+        room_nodes=room_nodes,
+        views_per_room=2,
+        habitat_sim_module=fake_habitat,
+        sensor_uuid="depth_sensor",
+        turn_action="turn_right",
+        progress_track=lambda iterable, description: iterable,
+    )
+
+    assert room_views == {"room_00": ["depth_0", "depth_1"]}
+    assert sim.actions == ["turn_right", "turn_right"]
+
+
+def test_collect_room_views_by_room_validates_views_per_room():
+    with pytest.raises(ValueError, match="views_per_room"):
+        collect_room_views_by_room(
+            agent=_FakeAgent(),
+            sim=_FakeSimulator(),
+            room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])],
+            views_per_room=0,
+            habitat_sim_module=SimpleNamespace(AgentState=_FakeAgentState),
+            progress_track=lambda iterable, description: iterable,
+        )