blackbird/pybot/tests/test_recorder.py

"""EpisodeRecorder 测试：actpred 格式数据集录制

全部在 dry_run / mock I2C 下运行，不碰真实硬件。
"""

import sys
from pathlib import Path

import numpy as np
import pytest
from PIL import Image

from pybot.recorder import EpisodeRecorder
from pybot.robot import RobotStruct

# ---- 测试用 .mt 脚本 ----

# 仅 @Camera + @Action(cam0.frame) —— 不需要 I2C
CAMERA_MT = """
@Camera
 name  width  height  framerate  capacity;
 cam0    320      240        15         2;
@

@Action
cam0.frame;
0;
@
"""

# @Camera + @Servo + 多行 @Action（含占位符）—— 需要 mock I2C
SERVO_CAMERA_MT = """
@Alias
name      alias;
   A      angle;
@

@I2C
 name        path;
i2c-1  /dev/i2c-1;
@

@PCA9685
     name    I2C   address;
pca9685-0  i2c-1   0x40;
@

@Servo
 name   max  min  offset  controller  channel;
  s1    90     -       -   pca9685-0        0;
  s2   180     -       -           -        1;
@

@Camera
 name  width  height  framerate  capacity;
 cam0    320      240        15         2;
@

@Action
s1.A  s2.A;
  45     0;
   -    10;
   -    20;
   -    30;
   -     0;
   -    10;
   -    20;
   -    30;
   -     0;
   -    10;
@
"""

# @Servo + @Action 但无 @Camera —— 用于测试无摄像头报错
NO_CAMERA_MT = """
@I2C
 name        path;
i2c-1  /dev/i2c-1;
@

@PCA9685
     name    I2C   address;
pca9685-0  i2c-1   0x40;
@

@Servo
 name   max  min  offset  controller  channel;
  s1    90     -       -   pca9685-0        0;
@

@Action
s1.A;
45;
@
"""


# ---- 测试 ----

def test_episode_format(tmp_path):
    """录制后目录结构正确（images/*.jpg + actions.npy）"""
    robot = RobotStruct(CAMERA_MT, dry_run=True)
    robot.active()
    rec = EpisodeRecorder(str(tmp_path), action_dim=18)
    robot.run(recorder=rec)
    robot.deactive()

    ep_dir = tmp_path / "episode_0000"
    assert ep_dir.is_dir()
    assert (ep_dir / "images").is_dir()
    assert (ep_dir / "actions.npy").exists()
    jpgs = list((ep_dir / "images").glob("*.jpg"))
    assert len(jpgs) == 1              # CAMERA_MT 有 1 行 @Action


def test_action_vector_padding(tmp_path, mock_i2c):
    """2 列 .mt → 18 维向量，末尾 16 维为零"""
    robot = RobotStruct(SERVO_CAMERA_MT, dry_run=True)
    robot.active()
    rec = EpisodeRecorder(str(tmp_path), action_dim=18)
    robot.run(recorder=rec)
    robot.deactive()

    actions = np.load(tmp_path / "episode_0000" / "actions.npy")
    assert actions.shape == (10, 18)    # 10 行 @Action × 18 维
    assert np.all(actions[:, 2:] == 0)  # 列 2-17 全为零


def test_placeholder_resolution(tmp_path, mock_i2c):
    """`-` 占位符正确递归解析"""
    robot = RobotStruct(SERVO_CAMERA_MT, dry_run=True)
    robot.active()
    rec = EpisodeRecorder(str(tmp_path), action_dim=18)
    robot.run(recorder=rec)
    robot.deactive()

    actions = np.load(tmp_path / "episode_0000" / "actions.npy")
    # s1.A: row0=45, row1..9 全是 '-' → 递归取上一行 → 全 45
    assert np.all(actions[:, 0] == 45)
    # s2.A: 0,10,20,30,0,10,20,30,0,10
    expected_s2 = [0, 10, 20, 30, 0, 10, 20, 30, 0, 10]
    assert list(actions[:, 1]) == expected_s2


def test_dry_run_synthetic_frame(tmp_path):
    """dry_run 下生成随机噪声帧，尺寸与 camera 配置一致"""
    robot = RobotStruct(CAMERA_MT, dry_run=True)
    robot.active()
    rec = EpisodeRecorder(str(tmp_path), action_dim=18)
    robot.run(recorder=rec)
    robot.deactive()

    img_path = tmp_path / "episode_0000" / "images" / "0000.jpg"
    assert img_path.exists()
    img = Image.open(img_path)
    assert img.size == (320, 240)      # (width, height)
    arr = np.array(img)
    assert arr.shape == (240, 320, 3)
    assert not np.all(arr == 0)          # 随机噪声，非全黑


def test_no_camera_error(tmp_path, mock_i2c):
    """非 dry_run 无 @Camera 时 --record 报错"""
    robot = RobotStruct(NO_CAMERA_MT, dry_run=False)
    robot.active()
    rec = EpisodeRecorder(str(tmp_path), action_dim=18)
    with pytest.raises(ValueError, match="--record 需要 @Camera"):
        robot.run(recorder=rec)
    robot.deactive()


def test_episode_auto_increment(tmp_path):
    """多次录制自动递增 episode 编号"""
    for _ in range(2):
        robot = RobotStruct(CAMERA_MT, dry_run=True)
        robot.active()
        rec = EpisodeRecorder(str(tmp_path), action_dim=18)
        robot.run(recorder=rec)
        robot.deactive()

    assert (tmp_path / "episode_0000").is_dir()
    assert (tmp_path / "episode_0001").is_dir()


def test_actpred_compatibility(tmp_path, mock_i2c):
    """用 actpred 的 ActionDataset 加载录制的 episode"""
    pytest.importorskip("torch")
    pytest.importorskip("torchvision")

    robot = RobotStruct(SERVO_CAMERA_MT, dry_run=True)
    robot.active()
    rec = EpisodeRecorder(str(tmp_path), action_dim=18)
    robot.run(recorder=rec)
    robot.deactive()

    actpred_dir = str(Path(__file__).resolve().parent.parent.parent / "actpred")
    sys.path.insert(0, actpred_dir)
    try:
        from dataset import ActionDataset
        ds = ActionDataset(str(tmp_path), history_len=5,
                           action_dim=18, frame_len=3)
        # 10 行，start=max(5,2)=5，samples: t=5..9 → 5 个样本
        assert len(ds) == 5
        sample = ds[0]
        # ActionDataset 默认 resize 到 640×480，与相机原始尺寸无关
        assert sample["images"].shape == (3, 3, 480, 640)  # (F, C, H, W)
        assert sample["action_history"].shape == (5, 18)
        assert sample["target"].shape == (18,)
    finally:
        sys.path.remove(actpred_dir)