tests/test_hardware.py

"""Tier 2 — Hardware integration tests for Marionette.

These tests require a physical Reachy Mini robot connected and powered on.
They exercise the full record/playback pipeline through the HTTP API with
real motor control.

All tests are decorated @pytest.mark.hardware and are skipped by default.
Run with:
    cd marionette
    pytest tests/test_hardware.py -m hardware

Or via the test runner:
    python tests/run_tests.py --hardware
"""

import json
import threading
import time
import wave
from pathlib import Path

import pytest
import uvicorn

from marionette.main import create_app, COUNTDOWN_SECONDS, MOTION_SAMPLE_RATE

# Skip all tests in this module unless -m hardware is specified
pytestmark = pytest.mark.hardware

HARDWARE_PORT = 18043
STARTUP_TIMEOUT = 60  # seconds to wait for robot startup animation
POLL_INTERVAL = 0.5


class _ServerThread:
    """Runs a uvicorn server in a background thread."""

    def __init__(self, app, port: int):
        self.config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
        self.server = uvicorn.Server(self.config)
        self.thread = threading.Thread(target=self.server.run, daemon=True)

    def start(self):
        self.thread.start()
        for _ in range(50):
            if self.server.started:
                break
            time.sleep(0.1)

    def stop(self):
        self.server.should_exit = True
        self.thread.join(timeout=5)


def _wait_for_mode(base_url: str, target_mode: str, timeout: float) -> dict:
    """Poll GET /api/state until mode matches target_mode or timeout."""
    import httpx

    deadline = time.time() + timeout
    last_state = {}
    while time.time() < deadline:
        try:
            resp = httpx.get(f"{base_url}/api/state", timeout=5)
            last_state = resp.json()
            if last_state.get("mode") == target_mode:
                return last_state
        except Exception:
            pass
        time.sleep(POLL_INTERVAL)
    raise TimeoutError(
        f"Timed out waiting for mode={target_mode!r} "
        f"(last mode={last_state.get('mode')!r})"
    )


@pytest.fixture(scope="session")
def hardware_server(tmp_path_factory):
    """Marionette server connected to a real robot.

    Starts the run() loop with a real ReachyMini and a uvicorn web server.
    Waits for the startup animation to complete (mode -> idle) before
    yielding the base URL and Marionette instance.
    """
    from reachy_mini import ReachyMini

    tmp = tmp_path_factory.mktemp("hardware")
    app, marionette = create_app(
        registry_path=tmp / "registry.json",
        dataset_root=tmp / "datasets",
    )

    reachy = ReachyMini()
    stop_event = threading.Event()

    # Start the run() loop — processes recordings/playbacks using the real robot
    run_thread = threading.Thread(
        target=marionette.run, args=(reachy, stop_event), daemon=True
    )
    run_thread.start()

    # Start the web server
    srv = _ServerThread(app, HARDWARE_PORT)
    srv.start()

    base_url = f"http://127.0.0.1:{HARDWARE_PORT}"

    # Wait for startup animation to finish
    try:
        _wait_for_mode(base_url, "idle", STARTUP_TIMEOUT)
    except TimeoutError:
        stop_event.set()
        srv.stop()
        raise

    yield base_url, marionette

    stop_event.set()
    srv.stop()


@pytest.fixture(scope="session")
def base_url(hardware_server):
    return hardware_server[0]


@pytest.fixture(scope="session")
def hw_marionette(hardware_server):
    return hardware_server[1]


# ──────── Tests ────────────────────────────────────────────────────────


class TestHardwareStartup:
    def test_startup_reaches_idle(self, base_url: str):
        """After fixture setup, verify mode is 'idle'."""
        import httpx

        resp = httpx.get(f"{base_url}/api/state", timeout=5)
        assert resp.status_code == 200
        assert resp.json()["mode"] == "idle"


class TestHardwareRecording:
    def test_record_captures_motion(self, base_url: str):
        """Record 2s silent, verify move JSON has ~200 frames at ~100Hz."""
        import httpx

        # Start recording
        resp = httpx.post(
            f"{base_url}/api/record",
            json={"duration": 2.0, "record_audio": False, "label": "hw-silent"},
            timeout=5,
        )
        assert resp.status_code == 200
        move_id = resp.json()["move_id"]

        # Wait for countdown + recording + saving to complete
        _wait_for_mode(base_url, "idle", timeout=15)

        # Verify move exists with expected frame count
        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        move = next((m for m in state["moves"] if m["id"] == move_id), None)
        assert move is not None, f"Move {move_id} not found in moves list"
        assert move["duration"] > 1.5, f"Duration too short: {move['duration']}"
        assert move["has_audio"] is False

    def test_record_with_audio(self, base_url: str):
        """Record 2s with mic, verify .wav file created."""
        import httpx

        resp = httpx.post(
            f"{base_url}/api/record",
            json={"duration": 2.0, "record_audio": True, "label": "hw-audio"},
            timeout=5,
        )
        assert resp.status_code == 200
        move_id = resp.json()["move_id"]

        _wait_for_mode(base_url, "idle", timeout=15)

        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        move = next((m for m in state["moves"] if m["id"] == move_id), None)
        assert move is not None, f"Move {move_id} not found"
        assert move["has_audio"] is True


class TestHardwarePlayback:
    def test_playback_silent_completes(self, base_url: str):
        """Play back a silent move, verify mode returns to idle."""
        import httpx

        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        move = next((m for m in state["moves"] if not m["has_audio"]), None)
        assert move is not None, "No silent move available for playback test"

        resp = httpx.post(
            f"{base_url}/api/play",
            json={"move_id": move["id"]},
            timeout=5,
        )
        assert resp.status_code == 200

        _wait_for_mode(base_url, "idle", timeout=move["duration"] + 15)

        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        assert state["mode"] == "idle"

    def test_playback_with_audio_completes(self, base_url: str):
        """Play back a move that has audio, verify sound path is exercised."""
        import httpx

        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        move = next((m for m in state["moves"] if m["has_audio"]), None)
        assert move is not None, (
            "No audio move available — test_record_with_audio must run first"
        )

        resp = httpx.post(
            f"{base_url}/api/play",
            json={"move_id": move["id"]},
            timeout=5,
        )
        assert resp.status_code == 200

        _wait_for_mode(base_url, "idle", timeout=move["duration"] + 15)

        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        assert state["mode"] == "idle"
        assert "finished" in state["message"].lower(), (
            f"Expected 'Finished playing' message, got: {state['message']}"
        )

    def test_record_and_delete(self, base_url: str):
        """Record a move, then delete it and verify file is removed."""
        import httpx

        # Record
        resp = httpx.post(
            f"{base_url}/api/record",
            json={"duration": 1.5, "record_audio": False, "label": "hw-delete-me"},
            timeout=5,
        )
        assert resp.status_code == 200
        move_id = resp.json()["move_id"]

        _wait_for_mode(base_url, "idle", timeout=15)

        # Verify it exists
        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        assert any(m["id"] == move_id for m in state["moves"])

        # Delete
        resp = httpx.delete(
            f"{base_url}/api/moves/{move_id}",
            timeout=5,
        )
        assert resp.status_code == 200

        # Verify it's gone
        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        assert not any(m["id"] == move_id for m in state["moves"])


class TestFullPipeline:
    """End-to-end pipeline tests: record → verify files → replay → delete.

    These tests exercise the full lifecycle and check that output files
    have the expected structure, durations, and frame counts.
    """

    def test_recording_produces_correct_json(self, base_url: str, hw_marionette):
        """Record 3s silent, verify JSON has timestamps, frames at ~100Hz."""
        import httpx

        duration = 3.0
        resp = httpx.post(
            f"{base_url}/api/record",
            json={"duration": duration, "record_audio": False, "label": "pipeline-json"},
            timeout=5,
        )
        assert resp.status_code == 200
        move_id = resp.json()["move_id"]

        _wait_for_mode(base_url, "idle", timeout=COUNTDOWN_SECONDS + duration + 10)

        # Read the actual JSON file from disk
        json_path = hw_marionette._dataset_dir / f"{move_id}.json"
        assert json_path.exists(), f"JSON file not found: {json_path}"

        data = json.loads(json_path.read_text())
        timestamps = data["time"]
        frames = data["set_target_data"]

        # Frame count should be close to duration * 100Hz
        expected_frames = int(duration * MOTION_SAMPLE_RATE)
        assert len(frames) > expected_frames * 0.8, (
            f"Too few frames: {len(frames)} (expected ~{expected_frames})"
        )
        assert len(timestamps) == len(frames)

        # Timestamps should span close to the requested duration
        actual_duration = timestamps[-1] - timestamps[0]
        assert actual_duration > duration * 0.8, (
            f"Recorded duration too short: {actual_duration:.2f}s (expected ~{duration}s)"
        )

        # Each frame must have head pose and antennas
        for frame in frames[:3]:  # check first few
            assert "head" in frame
            assert "antennas" in frame

        # Cleanup
        httpx.delete(f"{base_url}/api/moves/{move_id}", timeout=5)

    def test_recording_with_audio_produces_wav(self, base_url: str, hw_marionette):
        """Record 3s with mic, verify WAV file duration matches recording."""
        import httpx

        duration = 3.0
        resp = httpx.post(
            f"{base_url}/api/record",
            json={"duration": duration, "record_audio": True, "label": "pipeline-audio"},
            timeout=5,
        )
        assert resp.status_code == 200
        move_id = resp.json()["move_id"]

        _wait_for_mode(base_url, "idle", timeout=COUNTDOWN_SECONDS + duration + 10)

        # Verify WAV exists and has reasonable duration
        wav_path = hw_marionette._dataset_dir / f"{move_id}.wav"
        assert wav_path.exists(), f"WAV file not found: {wav_path}"

        with wave.open(str(wav_path), "rb") as wf:
            wav_duration = wf.getnframes() / wf.getframerate()
            assert wav_duration > duration * 0.7, (
                f"WAV too short: {wav_duration:.2f}s (expected ~{duration}s)"
            )
            assert wav_duration < duration + 2.0, (
                f"WAV too long: {wav_duration:.2f}s (expected ~{duration}s)"
            )

        # Also verify the JSON file reports has_audio
        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        move = next((m for m in state["moves"] if m["id"] == move_id), None)
        assert move is not None
        assert move["has_audio"] is True

        # Cleanup
        httpx.delete(f"{base_url}/api/moves/{move_id}", timeout=5)

    def test_record_replay_full_lifecycle(self, base_url: str, hw_marionette):
        """Record → verify → replay → verify playback completes → delete."""
        import httpx

        duration = 2.0

        # Step 1: Record
        resp = httpx.post(
            f"{base_url}/api/record",
            json={"duration": duration, "record_audio": False, "label": "lifecycle"},
            timeout=5,
        )
        assert resp.status_code == 200
        move_id = resp.json()["move_id"]

        state = _wait_for_mode(base_url, "idle", timeout=COUNTDOWN_SECONDS + duration + 10)
        assert any(m["id"] == move_id for m in state["moves"]), "Move not in list after recording"

        # Step 2: Replay
        resp = httpx.post(
            f"{base_url}/api/play",
            json={"move_id": move_id},
            timeout=5,
        )
        assert resp.status_code == 200

        state = _wait_for_mode(base_url, "idle", timeout=duration + 15)
        assert state["mode"] == "idle"

        # Step 3: Delete and verify cleanup
        json_path = hw_marionette._dataset_dir / f"{move_id}.json"
        assert json_path.exists()

        resp = httpx.delete(f"{base_url}/api/moves/{move_id}", timeout=5)
        assert resp.status_code == 200
        assert not json_path.exists(), "JSON file should be deleted"

        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        assert not any(m["id"] == move_id for m in state["moves"])

    def test_stop_cancels_queued_recording(self, base_url: str):
        """Submit a recording, immediately stop it — verify cancel works."""
        import httpx

        resp = httpx.post(
            f"{base_url}/api/record",
            json={"duration": 10.0, "record_audio": False, "label": "cancel-test"},
            timeout=5,
        )
        assert resp.status_code == 200

        # Immediately try to stop (may be queued or in countdown)
        time.sleep(0.1)
        resp = httpx.post(f"{base_url}/api/record/stop", timeout=5)
        assert resp.status_code == 200
        data = resp.json()
        assert data["stopped"] is True

        # Should return to idle
        _wait_for_mode(base_url, "idle", timeout=10)

    def test_recording_transitions_through_phases(self, base_url: str):
        """Verify the recording goes through queued → countdown → recording → idle."""
        import httpx

        duration = 3.0
        observed_modes = set()

        resp = httpx.post(
            f"{base_url}/api/record",
            json={"duration": duration, "record_audio": False, "label": "phases"},
            timeout=5,
        )
        assert resp.status_code == 200

        # Poll rapidly to observe phase transitions
        deadline = time.time() + COUNTDOWN_SECONDS + duration + 10
        while time.time() < deadline:
            try:
                state = httpx.get(f"{base_url}/api/state", timeout=2).json()
                observed_modes.add(state["mode"])

                # When in recording mode, verify timing fields are present
                if state["mode"] == "recording":
                    assert state["phase_start_at"] is not None, "recording phase_start_at is null"
                    assert state["phase_end_at"] is not None, "recording phase_end_at is null"

                # When in countdown, verify timing fields
                if state["mode"] == "countdown":
                    assert state["phase_start_at"] is not None, "countdown phase_start_at is null"
                    assert state["phase_end_at"] is not None, "countdown phase_end_at is null"

                if state["mode"] == "idle" and "phases" in (state.get("message") or ""):
                    break  # Recording completed
            except Exception:
                pass
            time.sleep(0.15)

        # We should have seen at least countdown and recording phases
        assert "countdown" in observed_modes, (
            f"Never saw countdown mode. Observed: {observed_modes}"
        )
        assert "recording" in observed_modes, (
            f"Never saw recording mode. Observed: {observed_modes}"
        )
        assert "idle" in observed_modes, (
            f"Never returned to idle. Observed: {observed_modes}"
        )

        # Cleanup
        state = httpx.get(f"{base_url}/api/state", timeout=5).json()
        move = next((m for m in state["moves"] if "phases" in m.get("label", "")), None)
        if move:
            httpx.delete(f"{base_url}/api/moves/{move['id']}", timeout=5)