"""Online robot preprocessing checks shared by Turtle serving."""

from types import SimpleNamespace

import numpy as np

from wall_x._vendor.harrix.serving._wallx_infer.robot import Robot


def test_dof_mask_disables_virtual_action_padding():
    wrapper = SimpleNamespace(
        config=SimpleNamespace(
            action_horizon=10,
            train_config={
                "dof_config": {
                    "master_right_ee_cartesian_pos": 3,
                    "master_right_ee_rotation": 3,
                    "master_right_gripper": 1,
                    "action_padding": 19,
                }
            },
        )
    )

    mask = Robot._get_dof_mask(wrapper)

    assert mask.shape == (1, 10, 26)
    np.testing.assert_array_equal(mask[:, :, :7], 1)
    np.testing.assert_array_equal(mask[:, :, 7:], 0)