434 lines
11 KiB
Python
434 lines
11 KiB
Python
from typing import List, Optional, Type, Union
|
|
|
|
import pytest
|
|
|
|
from sampy.raknet.bitstream import Bitstream, get_bit_length
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"value",
|
|
[
|
|
b"",
|
|
b"A",
|
|
b"AB",
|
|
b"\xFF\x00\xAA\x55",
|
|
],
|
|
)
|
|
def test_from_bytes(value: bytes):
|
|
bitstream = Bitstream.from_bytes(value)
|
|
assert len(bitstream) == (len(value) << 3)
|
|
assert bytes(bitstream) == value
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"value,length,expected",
|
|
[
|
|
(0b0, 0, None),
|
|
(0b1, 1, None),
|
|
(0b0, 1, None),
|
|
(0b1011, 4, None),
|
|
(0b10111011, 8, None),
|
|
(0b10111011, 7, 0b0111011), # -> 0b0111011
|
|
],
|
|
)
|
|
def test_from_int(value: int, length: int, expected: int):
|
|
bitstream = Bitstream.from_int(value, length)
|
|
assert len(bitstream) == length
|
|
assert len(bytes(bitstream)) == ((length + 7) >> 3)
|
|
assert bitstream.bits == [bool(value & (1 << i)) for i in range(length - 1, -1, -1)]
|
|
if expected is not None:
|
|
expected_bitstream = Bitstream.from_int(expected, length)
|
|
assert bitstream.bits == expected_bitstream.bits
|
|
assert bytes(bitstream) == bytes(expected_bitstream)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"values",
|
|
[
|
|
[True, False, False, True, True],
|
|
[[True], [False, True], [False, False]],
|
|
[b"a", b"bc"],
|
|
[b"aa", b"b"],
|
|
[
|
|
Bitstream.from_bytes(b"A"),
|
|
Bitstream.from_int(0b1011, 4),
|
|
Bitstream.from_bytes(b"B"),
|
|
],
|
|
[True, b"A", b"B", [False, True], Bitstream.from_int(0b10110, 5), False],
|
|
],
|
|
)
|
|
def test_init(values: List[Union[bool, List[bool], bytes, Bitstream]]):
|
|
bitstream = Bitstream(*values)
|
|
|
|
bitstream.seek(0)
|
|
for item in values:
|
|
bit_length = get_bit_length(item)
|
|
if type(item) is bool:
|
|
item = [item]
|
|
|
|
read_type = bool if type(item) is list else type(item)
|
|
value = bitstream.read(read_type, bit_length)
|
|
assert value == item
|
|
|
|
|
|
def test_getitem():
|
|
sample_values = [True, False, False, True, False, True, True, True, True, False]
|
|
|
|
for length in range(1, len(sample_values)):
|
|
values = sample_values[:length]
|
|
bitstream = Bitstream(values)
|
|
for i in range(length):
|
|
assert bitstream[i] == values[i]
|
|
assert bitstream[-i - 1] == values[-i - 1] # For inverse lookup
|
|
|
|
|
|
def test_getitem_index_error():
|
|
bitstream = Bitstream(True)
|
|
with pytest.raises(IndexError):
|
|
bitstream[1]
|
|
with pytest.raises(IndexError):
|
|
bitstream[-2] # Inverse IndexError
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"value",
|
|
[
|
|
True,
|
|
[True, True],
|
|
b"A",
|
|
Bitstream(True, b"C"),
|
|
],
|
|
)
|
|
def test_setitem(value: Union[bool, List[bool], bytes, Bitstream]):
|
|
bit_length = get_bit_length(value)
|
|
read_type = bool if type(value) is list else type(value)
|
|
for length in range(10):
|
|
for i in range(length - bit_length + 1):
|
|
bitstream = Bitstream([False] * length)
|
|
bitstream[i] = value
|
|
assert bitstream.read(read_type, bit_length, i) == (
|
|
value if type(value) is not bool else [value]
|
|
)
|
|
assert len(bitstream) == length
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"value,equals,not_equals",
|
|
[
|
|
([], [[]], [[True], [False]]),
|
|
(True, [[True]], [[False], [True, False], b"A"]),
|
|
([True, False], [[True, False]], [[False], [False, False], b"A"]),
|
|
(
|
|
b"A",
|
|
[b"A", b"A", Bitstream.from_bytes(b"A"), Bitstream.from_bytes(b"A").bits],
|
|
[
|
|
[],
|
|
[False],
|
|
b"B",
|
|
b"B",
|
|
Bitstream.from_bytes(b"B"),
|
|
Bitstream.from_bytes(b"B").bits,
|
|
],
|
|
),
|
|
(b"ABC", [b"ABC"], [[], [False], b"A"]),
|
|
],
|
|
)
|
|
def test_eq(
|
|
value: Union[bool, List[bool], bytes, Bitstream],
|
|
equals: List[Union[bool, List[bool], bytes, Bitstream]],
|
|
not_equals: List[Union[bool, List[bool], bytes, Bitstream]],
|
|
):
|
|
bitstream = Bitstream(value)
|
|
for equal in equals:
|
|
assert bitstream == equal
|
|
for not_equal in not_equals:
|
|
assert bitstream != not_equal
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"values,number,expected",
|
|
[
|
|
([], 2, []),
|
|
([True, False], 2, [True, False, True, False]),
|
|
(b"A", 2, b"AA"),
|
|
([False], 3, [False, False, False]),
|
|
],
|
|
)
|
|
def test_mul(
|
|
values: Union[bool, List[bool], bytes, Bitstream],
|
|
number: int,
|
|
expected: Union[List[bool], bytes, Bitstream],
|
|
):
|
|
bitstream = Bitstream(values)
|
|
|
|
assert (bitstream * number) == expected
|
|
assert bitstream == values
|
|
|
|
bitstream *= number
|
|
assert bitstream == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"init_values,add_value,expected",
|
|
[
|
|
([], True, [True]),
|
|
([True, False], False, [True, False, False]),
|
|
(b"A", b"B", b"AB"),
|
|
],
|
|
)
|
|
def test_add(
|
|
init_values: Union[bool, List[bool], bytes, Bitstream],
|
|
add_value: Union[bool, List[bool], bytes, Bitstream],
|
|
expected: Union[List[bool], bytes, Bitstream],
|
|
):
|
|
bitstream = Bitstream(init_values)
|
|
|
|
assert (bitstream + add_value) == expected
|
|
assert len(bitstream) == get_bit_length(init_values)
|
|
|
|
bitstream += add_value
|
|
assert bitstream == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"init_values",
|
|
[
|
|
[],
|
|
[True, False],
|
|
[b"A"],
|
|
[True, b"B"],
|
|
],
|
|
)
|
|
def test_iter(init_values: List[Union[bool, List[bool], bytes, Bitstream]]):
|
|
bitstream = Bitstream(*init_values)
|
|
assert len(list(bitstream)) == len(bitstream)
|
|
for index, bit in enumerate(bitstream):
|
|
assert bit == bitstream[index]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"init_values,expected",
|
|
[
|
|
([], b""),
|
|
([b"A"], b"A"),
|
|
([False], b"\x00"),
|
|
([[False] * 6, True], b"\x02"),
|
|
([[False] * 7, True], b"\x01"),
|
|
],
|
|
)
|
|
def test_bytes(
|
|
init_values: List[Union[bool, List[bool], bytes, Bitstream]], expected: bytes
|
|
):
|
|
assert bytes(Bitstream(*init_values)) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"init_values,expected",
|
|
[
|
|
([], []),
|
|
([b"\x01"], [False] * 7 + [True]),
|
|
([False], [False]),
|
|
([False] * 6 + [True], [False] * 6 + [True]),
|
|
([False] * 7 + [True], [False] * 7 + [True]),
|
|
],
|
|
)
|
|
def test_bits(
|
|
init_values: List[Union[bool, List[bool], bytes, Bitstream]], expected: List[bool]
|
|
):
|
|
assert Bitstream(*init_values).bits == expected
|
|
|
|
|
|
def test_copy():
|
|
bitstream = Bitstream(True, False)
|
|
bitstream_copy = bitstream.copy()
|
|
assert bitstream == bitstream_copy
|
|
assert id(bitstream) != id(bitstream.copy)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"lst",
|
|
[
|
|
[False] * 10,
|
|
[True] * 10,
|
|
[True, False] * 5,
|
|
[False, True] * 5,
|
|
],
|
|
)
|
|
def test_append(lst: List[bool]):
|
|
bitstream = Bitstream()
|
|
for i in range(len(lst)):
|
|
bitstream.append(lst[i])
|
|
assert bitstream[i] == lst[i]
|
|
assert len(bitstream) == (i + 1)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"extend_value",
|
|
[
|
|
[True, False, True],
|
|
b"A",
|
|
b"BB",
|
|
Bitstream(True, b"C"),
|
|
],
|
|
)
|
|
def test_extend(extend_value: Union[List[bool], bytes, Bitstream]):
|
|
for i in range(10):
|
|
bitstream = Bitstream([False] * i)
|
|
bitstream.extend(extend_value)
|
|
value = bitstream.read(
|
|
bool if type(extend_value) is list else type(extend_value),
|
|
get_bit_length(extend_value),
|
|
i,
|
|
)
|
|
assert value == extend_value
|
|
|
|
|
|
def test_extend_error():
|
|
bitstream = Bitstream()
|
|
with pytest.raises(TypeError):
|
|
bitstream.extend(True)
|
|
|
|
|
|
def test_clear():
|
|
bitstream = Bitstream(b"A", [True, False, False, True])
|
|
bitstream._offset = 3
|
|
bitstream.clear()
|
|
assert len(bitstream) == 0
|
|
assert len(bytes(bitstream)) == 0
|
|
assert bitstream._offset == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"lst",
|
|
[
|
|
[False] * 10,
|
|
[True] * 10,
|
|
[True, False] * 5,
|
|
[False, True] * 5,
|
|
],
|
|
)
|
|
def test_pop(lst: List[bool]):
|
|
bitstream = Bitstream(lst)
|
|
for i in range(len(bitstream)):
|
|
assert bitstream.pop() == lst[-i - 1]
|
|
|
|
with pytest.raises(IndexError):
|
|
assert bitstream.pop()
|
|
|
|
|
|
def test_read():
|
|
bitstream = Bitstream(True, b"A", False, True, False)
|
|
|
|
assert bitstream.read(bool, 1) == [True]
|
|
assert bitstream.read(bytes, 8) == b"A"
|
|
assert bitstream.read(int, 3) == 2
|
|
|
|
bitstream.seek(0)
|
|
assert bitstream.read(Bitstream, len(bitstream) - 1) == Bitstream(
|
|
True, b"A", False, True
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"bitstream,read_type,bit_length,bit_index,expected",
|
|
[
|
|
(Bitstream(True, b"A", False, True, False), bool, 1, 0, [True]),
|
|
(
|
|
Bitstream(True, b"A", False, True, False),
|
|
bool,
|
|
3,
|
|
8 + 1,
|
|
[False, True, False],
|
|
),
|
|
(Bitstream(True, b"A", False, True, False), bytes, 8, 1, b"A"),
|
|
(
|
|
Bitstream(True, b"A", False, True, False),
|
|
Bitstream,
|
|
1 + 8 + 3,
|
|
0,
|
|
Bitstream(True, b"A", False, True, False),
|
|
),
|
|
(Bitstream(True, b"A", False, True, False), int, 3, 1 + 8, 2),
|
|
],
|
|
)
|
|
def test_read_index(
|
|
bitstream: Bitstream,
|
|
read_type: Type[Union[bool, bytes, Bitstream, int]],
|
|
bit_length: int,
|
|
bit_index: Optional[int],
|
|
expected: Union[List[bool], bytes, Bitstream, int],
|
|
):
|
|
value = bitstream.read(read_type, bit_length, bit_index)
|
|
|
|
if read_type is bool:
|
|
assert type(value) == list
|
|
else:
|
|
assert type(value) == read_type
|
|
assert value == expected
|
|
|
|
|
|
def test_read_error():
|
|
bitstream = Bitstream(True, b"A", False, True, False)
|
|
|
|
with pytest.raises(IndexError):
|
|
bitstream.read(bool, 1, len(bitstream))
|
|
|
|
with pytest.raises(IndexError):
|
|
bitstream.read(bool, 1, -1)
|
|
|
|
with pytest.raises(TypeError):
|
|
bitstream.read(list, 1, 0)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"bitstream,value,bit_index,current_offset",
|
|
[
|
|
(Bitstream(), True, None, None),
|
|
(Bitstream(), b"A", None, None),
|
|
(Bitstream(), True, 0, None),
|
|
(Bitstream(), True, 3, None),
|
|
(Bitstream(), True, 7, None),
|
|
(Bitstream(), True, 8, None),
|
|
(Bitstream(False), True, None, None),
|
|
(Bitstream(False), True, 0, None),
|
|
(Bitstream(False), True, 0, 0),
|
|
(Bitstream(False), True, None, 0),
|
|
],
|
|
)
|
|
def test_write(
|
|
bitstream: Bitstream,
|
|
value: Union[bool, List[bool], bytes, Bitstream],
|
|
bit_index: Optional[int],
|
|
current_offset: Optional[int],
|
|
):
|
|
if current_offset is not None:
|
|
bitstream.seek(current_offset)
|
|
|
|
bit_length = get_bit_length(value)
|
|
bitstream.write(value, bit_index)
|
|
|
|
if bit_index is None:
|
|
bitstream.seek(-bit_length, 1)
|
|
read_value = bitstream.read(
|
|
bool if type(value) is list else type(value), bit_length
|
|
)
|
|
else:
|
|
read_value = bitstream.read(
|
|
bool if type(value) is list else type(value), bit_length, bit_index
|
|
)
|
|
|
|
assert read_value == ([value] if type(value) is bool else value)
|
|
|
|
|
|
def test_seek():
|
|
bitstream = Bitstream([True] * 8)
|
|
|
|
bitstream.seek(4, 0)
|
|
assert bitstream._offset == 4
|
|
|
|
bitstream.seek(-1, 1)
|
|
assert bitstream._offset == 3
|
|
|
|
bitstream.seek(-1, 2)
|
|
assert bitstream._offset == 7
|