Mini-Nav/tests/test_habitat_depth_sensor.py

"""Tests for optional Habitat depth sensor support in create_habitat_simulator."""

from __future__ import annotations

import sys
from pathlib import Path
from types import SimpleNamespace

import pytest

MINI_NAV_DIR = Path(__file__).resolve().parents[1] / "mini-nav"
sys.path.insert(0, str(MINI_NAV_DIR))

from simulator.habitat import (  # noqa: E402
    HabitatSimulatorConfig,
    create_habitat_simulator,
)

# ---------------------------------------------------------------------------
# Fake habitat_sim module
# ---------------------------------------------------------------------------

_specs_created: list = []


class _FakeCameraSensorSpec:
    """Fake CameraSensorSpec that records itself in _specs_created."""

    def __init__(self) -> None:
        self.uuid = None
        self.sensor_type = None
        self.resolution = None
        self.position = None
        self.hfov = None
        _specs_created.append(self)


class _FakeSensorType:
    COLOR = "COLOR"
    DEPTH = "DEPTH"


class _FakeSimulator:
    def __init__(self, cfg: object) -> None:
        self.cfg = cfg

    @staticmethod
    def initialize_agent(agent_id: int) -> SimpleNamespace:
        return SimpleNamespace(agent_id=agent_id)

    @staticmethod
    def close() -> None:
        return None


def _make_fake_habitat_sim_module() -> SimpleNamespace:
    """Return a fake ``habitat_sim`` module for testing."""
    global _specs_created
    _specs_created = []

    hab_sim = SimpleNamespace(
        SimulatorConfiguration=SimpleNamespace,
        CameraSensorSpec=_FakeCameraSensorSpec,
        SensorType=_FakeSensorType,
        Configuration=lambda sim_cfg, agent_cfgs: SimpleNamespace(
            sim_cfg=sim_cfg,
            agent_cfgs=agent_cfgs,
        ),
        Simulator=_FakeSimulator,
    )

    class _FakeActionSpec:
        def __init__(self, name: str, actuation: object) -> None:
            self.name = name
            self.actuation = actuation

    class _FakeActuationSpec:
        def __init__(self, amount: float | None = None) -> None:
            self.amount = amount

    hab_sim.agent = SimpleNamespace(
        AgentConfiguration=SimpleNamespace,
        ActionSpec=_FakeActionSpec,
        ActuationSpec=_FakeActuationSpec,
    )

    return hab_sim


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def fake_habitat_sim() -> SimpleNamespace:
    """Fixture providing a fresh fake habitat_sim module per test."""
    return _make_fake_habitat_sim_module()


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


def test_default_config_creates_color_and_depth_sensor_specs(
    fake_habitat_sim: SimpleNamespace,
) -> None:
    """Default config should produce two specs with expected UUIDs and types."""
    config = HabitatSimulatorConfig(scene_path="scene.glb")
    create_habitat_simulator(config, fake_habitat_sim)

    assert len(_specs_created) == 2
    assert [s.uuid for s in _specs_created] == ["color_sensor", "depth_sensor"]
    assert [s.sensor_type for s in _specs_created] == ["COLOR", "DEPTH"]


def test_create_habitat_simulator_adds_matching_rgb_and_depth_specs(
    fake_habitat_sim: SimpleNamespace,
) -> None:
    """Depth sensor spec should mirror the RGB spec when depth_sensor_uuid is set."""
    config = HabitatSimulatorConfig(
        scene_path="scene.glb",
        image_size=320,
        sensor_height=1.25,
        sensor_uuid="rgb",
        depth_sensor_uuid="depth",
        hfov_degrees=75.0,
    )

    create_habitat_simulator(config, fake_habitat_sim)

    assert len(_specs_created) == 2, "Expected 2 sensor specs (RGB + depth)"

    rgb_spec, depth_spec = _specs_created

    # UUIDs
    assert rgb_spec.uuid == "rgb"
    assert depth_spec.uuid == "depth"

    # Sensor types
    assert rgb_spec.sensor_type == "COLOR"
    assert depth_spec.sensor_type == "DEPTH"

    # Resolutions
    assert rgb_spec.resolution == [320, 320]
    assert depth_spec.resolution == [320, 320]

    # Positions
    assert rgb_spec.position == [0.0, 1.25, 0.0]
    assert depth_spec.position == [0.0, 1.25, 0.0]

    # HFOV
    assert rgb_spec.hfov == 75.0
    assert depth_spec.hfov == 75.0


def test_create_habitat_simulator_can_disable_depth_sensor(
    fake_habitat_sim: SimpleNamespace,
) -> None:
    """``depth_sensor_uuid=None`` should produce only a single color sensor spec."""
    config = HabitatSimulatorConfig(
        scene_path="scene.glb",
        depth_sensor_uuid=None,
    )

    create_habitat_simulator(config, fake_habitat_sim)

    assert len(_specs_created) == 1
    assert _specs_created[0].uuid == "color_sensor"


def test_habitat_config_rejects_invalid_hfov(
    fake_habitat_sim: SimpleNamespace,
) -> None:
    """``hfov_degrees=0`` should raise ``ValueError`` containing ``hfov_degrees``."""
    config = HabitatSimulatorConfig(
        scene_path="scene.glb",
        hfov_degrees=0.0,
    )

    with pytest.raises(ValueError, match="hfov_degrees"):
        create_habitat_simulator(config, fake_habitat_sim)