diff --git a/chia/util/struct_stream.py b/chia/util/struct_stream.py index 5ac9436470..726eba0efc 100644 --- a/chia/util/struct_stream.py +++ b/chia/util/struct_stream.py @@ -87,6 +87,5 @@ class StructStream(int): def stream_to_bytes(self) -> bytes: return super().to_bytes(length=self.SIZE, byteorder="big", signed=self.SIGNED) - # this is meant to avoid mixing up construcing a bytes object of a specific - # size (i.e. bytes(int)) vs. serializing the integer to bytes (i.e. bytes(uint32)) - __bytes__ = None + def __bytes__(self) -> bytes: + return self.stream_to_bytes() diff --git a/tests/util/test_struct_stream.py b/tests/util/test_struct_stream.py index f07a46d3e6..37ce670ab7 100644 --- a/tests/util/test_struct_stream.py +++ b/tests/util/test_struct_stream.py @@ -1,5 +1,6 @@ from __future__ import annotations +import enum import io import struct from dataclasses import dataclass @@ -27,6 +28,13 @@ def dataclass_parameters(instances: Iterable[object]) -> List[ParameterSet]: return [dataclass_parameter(instance) for instance in instances] +class StreamAndBytesMatchMode(enum.Enum): + minimum = "minimum" + middle_low = "middle low" + middle_high = "middle high" + maximum = "maximum" + + @dataclass(frozen=True) class BadName: name: str @@ -288,3 +296,19 @@ class TestStructStream: def test_parse_metadata_from_name_correct_minimum(self, good: Good) -> None: assert good.cls.MINIMUM == good.minimum + + @pytest.mark.parametrize("mode", list(StreamAndBytesMatchMode), ids=lambda mode: mode.value) + def test_stream_to_bytes_and_bytes_match_minimum(self, good: Good, mode: StreamAndBytesMatchMode) -> None: + if mode == StreamAndBytesMatchMode.minimum: + value = good.minimum + elif mode == StreamAndBytesMatchMode.middle_low: + value = int(good.minimum + ((good.maximum - good.minimum) * 0.3)) + elif mode == StreamAndBytesMatchMode.middle_high: + value = int(good.minimum + ((good.maximum - good.minimum) * 0.7)) + elif mode == StreamAndBytesMatchMode.maximum: + value = good.maximum + else: + raise Exception(f"unhandled parametrization: {mode!r}") + + instance = good.cls(value) + assert bytes(instance) == instance.stream_to_bytes()