mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
109 lines
3.0 KiB
Python
109 lines
3.0 KiB
Python
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,
|
|
)
|