lerobot/tests/utils/test_encoding_utils.py

156 lines
3.7 KiB
Python

import pytest
from lerobot.common.utils.encoding_utils import (
decode_sign_magnitude,
decode_twos_complement,
encode_sign_magnitude,
encode_twos_complement,
)
@pytest.mark.parametrize(
"value, sign_bit_index, expected",
[
(5, 4, 5),
(0, 4, 0),
(7, 3, 7),
(-1, 4, 17),
(-8, 4, 24),
(-3, 3, 11),
],
)
def test_encode_sign_magnitude(value, sign_bit_index, expected):
assert encode_sign_magnitude(value, sign_bit_index) == expected
@pytest.mark.parametrize(
"encoded, sign_bit_index, expected",
[
(5, 4, 5),
(0, 4, 0),
(7, 3, 7),
(17, 4, -1),
(24, 4, -8),
(11, 3, -3),
],
)
def test_decode_sign_magnitude(encoded, sign_bit_index, expected):
assert decode_sign_magnitude(encoded, sign_bit_index) == expected
@pytest.mark.parametrize(
"encoded, sign_bit_index",
[
(16, 4),
(-9, 3),
],
)
def test_encode_raises_on_overflow(encoded, sign_bit_index):
with pytest.raises(ValueError):
encode_sign_magnitude(encoded, sign_bit_index)
def test_encode_decode_sign_magnitude():
for sign_bit_index in range(2, 6):
max_val = (1 << sign_bit_index) - 1
for value in range(-max_val, max_val + 1):
encoded = encode_sign_magnitude(value, sign_bit_index)
decoded = decode_sign_magnitude(encoded, sign_bit_index)
assert decoded == value, f"Failed at value={value}, index={sign_bit_index}"
@pytest.mark.parametrize(
"value, n_bytes, expected",
[
(0, 1, 0),
(5, 1, 5),
(-1, 1, 255),
(-128, 1, 128),
(-2, 1, 254),
(127, 1, 127),
(0, 2, 0),
(5, 2, 5),
(-1, 2, 65_535),
(-32_768, 2, 32_768),
(-2, 2, 65_534),
(32_767, 2, 32_767),
(0, 4, 0),
(5, 4, 5),
(-1, 4, 4_294_967_295),
(-2_147_483_648, 4, 2_147_483_648),
(-2, 4, 4_294_967_294),
(2_147_483_647, 4, 2_147_483_647),
],
)
def test_encode_twos_complement(value, n_bytes, expected):
assert encode_twos_complement(value, n_bytes) == expected
@pytest.mark.parametrize(
"value, n_bytes, expected",
[
(0, 1, 0),
(5, 1, 5),
(255, 1, -1),
(128, 1, -128),
(254, 1, -2),
(127, 1, 127),
(0, 2, 0),
(5, 2, 5),
(65_535, 2, -1),
(32_768, 2, -32_768),
(65_534, 2, -2),
(32_767, 2, 32_767),
(0, 4, 0),
(5, 4, 5),
(4_294_967_295, 4, -1),
(2_147_483_648, 4, -2_147_483_648),
(4_294_967_294, 4, -2),
(2_147_483_647, 4, 2_147_483_647),
],
)
def test_decode_twos_complement(value, n_bytes, expected):
assert decode_twos_complement(value, n_bytes) == expected
@pytest.mark.parametrize(
"value, n_bytes",
[
(-129, 1),
(128, 1),
(-32_769, 2),
(32_768, 2),
(-2_147_483_649, 4),
(2_147_483_648, 4),
],
)
def test_encode_twos_complement_out_of_range(value, n_bytes):
with pytest.raises(ValueError):
encode_twos_complement(value, n_bytes)
@pytest.mark.parametrize(
"value, n_bytes",
[
(-128, 1),
(-1, 1),
(0, 1),
(1, 1),
(127, 1),
(-32_768, 2),
(-1, 2),
(0, 2),
(1, 2),
(32_767, 2),
(-2_147_483_648, 4),
(-1, 4),
(0, 4),
(1, 4),
(2_147_483_647, 4),
],
)
def test_encode_decode_twos_complement(value, n_bytes):
encoded = encode_twos_complement(value, n_bytes)
decoded = decode_twos_complement(encoded, n_bytes)
assert decoded == value, f"Failed at value={value}, n_bytes={n_bytes}"