FusionAGI/tests/test_embodiment.py

"""Tests for Embodied Intelligence / Robotics bridge."""

from __future__ import annotations

import pytest

from fusionagi.maa.embodiment import (
    ActuatorState,
    EmbodimentBridge,
    MotionCommand,
    SensorType,
    SimulatedActuator,
    SimulatedSensor,
    TrajectoryPoint,
)


@pytest.fixture
def actuator() -> SimulatedActuator:
    return SimulatedActuator(joint_ids=["j0", "j1", "j2"])


@pytest.fixture
def sensor() -> SimulatedSensor:
    s = SimulatedSensor()
    s.register_sensor("cam0", SensorType.CAMERA, {"width": 640, "height": 480})
    s.register_sensor("imu0", SensorType.IMU, {"accel": [0, 0, 9.8]})
    return s


@pytest.fixture
def bridge(actuator: SimulatedActuator, sensor: SimulatedSensor) -> EmbodimentBridge:
    return EmbodimentBridge(actuator=actuator, sensors=sensor)


class TestSimulatedActuator:
    @pytest.mark.asyncio
    async def test_get_joint_states(self, actuator: SimulatedActuator) -> None:
        states = await actuator.get_joint_states()
        assert len(states) == 3
        assert all(s.position == 0.0 for s in states)

    @pytest.mark.asyncio
    async def test_get_state_idle(self, actuator: SimulatedActuator) -> None:
        state = await actuator.get_state()
        assert state == ActuatorState.IDLE

    @pytest.mark.asyncio
    async def test_execute_motion(self, actuator: SimulatedActuator) -> None:
        cmd = MotionCommand(
            command_id="test_cmd",
            trajectory=[
                TrajectoryPoint(joint_positions={"j0": 1.0, "j1": -0.5}, time_from_start=1.0)
            ],
        )
        result = await actuator.execute_motion(cmd)
        assert result.success
        assert result.command_id == "test_cmd"

    @pytest.mark.asyncio
    async def test_emergency_stop(self, actuator: SimulatedActuator) -> None:
        assert await actuator.emergency_stop()
        state = await actuator.get_state()
        assert state == ActuatorState.EMERGENCY_STOP


class TestSimulatedSensor:
    @pytest.mark.asyncio
    async def test_list_sensors(self, sensor: SimulatedSensor) -> None:
        ids = await sensor.list_sensors()
        assert "cam0" in ids
        assert "imu0" in ids

    @pytest.mark.asyncio
    async def test_read_sensor(self, sensor: SimulatedSensor) -> None:
        reading = await sensor.read("cam0")
        assert reading is not None
        assert reading.sensor_type == SensorType.CAMERA

    @pytest.mark.asyncio
    async def test_read_missing_sensor(self, sensor: SimulatedSensor) -> None:
        reading = await sensor.read("nonexistent")
        assert reading is None


class TestEmbodimentBridge:
    @pytest.mark.asyncio
    async def test_perceive(self, bridge: EmbodimentBridge) -> None:
        perception = await bridge.perceive()
        assert "sensors" in perception
        assert "joints" in perception
        assert len(perception["joints"]) == 3

    @pytest.mark.asyncio
    async def test_execute_within_bounds(self, bridge: EmbodimentBridge) -> None:
        cmd = MotionCommand(
            command_id="cmd1",
            trajectory=[TrajectoryPoint(joint_positions={"j0": 0.5}, time_from_start=1.0)],
        )
        result = await bridge.execute(cmd)
        assert result.success

    @pytest.mark.asyncio
    async def test_execute_workspace_bounds_violated(self) -> None:
        actuator = SimulatedActuator(joint_ids=["j0"])
        bridge = EmbodimentBridge(
            actuator=actuator,
            workspace_bounds={"j0": (-1.0, 1.0)},
        )
        cmd = MotionCommand(
            command_id="cmd_oob",
            trajectory=[TrajectoryPoint(joint_positions={"j0": 5.0}, time_from_start=1.0)],
        )
        result = await bridge.execute(cmd)
        assert not result.success
        assert "outside bounds" in result.error_message

    @pytest.mark.asyncio
    async def test_execute_no_actuator(self) -> None:
        bridge = EmbodimentBridge()
        cmd = MotionCommand(command_id="cmd_none", trajectory=[])
        result = await bridge.execute(cmd)
        assert not result.success

    @pytest.mark.asyncio
    async def test_stop(self, bridge: EmbodimentBridge) -> None:
        assert await bridge.stop()

    def test_get_summary(self, bridge: EmbodimentBridge) -> None:
        summary = bridge.get_summary()
        assert summary["actuator_connected"]
        assert summary["sensors_connected"]