diff --git a/py/farm_ng/core/events_file_writer.py b/py/farm_ng/core/events_file_writer.py index 682eaa4e..39746d38 100644 --- a/py/farm_ng/core/events_file_writer.py +++ b/py/farm_ng/core/events_file_writer.py @@ -7,6 +7,7 @@ # pylint can't find Event or Uri in protobuf generated files # https://github.com/protocolbuffers/protobuf/issues/10372 from farm_ng.core.event_pb2 import Event +from farm_ng.core.events_file_reader import payload_to_protobuf from farm_ng.core.stamp import StampSemantics, get_monotonic_now, get_system_clock_now from farm_ng.core.uri import make_proto_uri, uri_to_string from google.protobuf.json_format import MessageToJson @@ -152,6 +153,18 @@ def header_events(self) -> dict[str, tuple[Event, bytes]]: """ return self._header_events + @property + def header_messages(self) -> list[Message]: + """Return the header_events, formatted as a list of protobuf messages. + + Returns: + list[Message]: List of header messages. + """ + return [ + payload_to_protobuf(event, payload) + for (event, payload) in self.header_events.values() + ] + def add_header_event( self, event: Event, diff --git a/py/tests/test_events_file_writer.py b/py/tests/test_events_file_writer.py index 643b185d..2b18a05b 100644 --- a/py/tests/test_events_file_writer.py +++ b/py/tests/test_events_file_writer.py @@ -13,6 +13,7 @@ from pathlib import Path from farm_ng.core.uri_pb2 import Uri + from google.protobuf.message import Message @pytest.fixture() @@ -132,6 +133,7 @@ def test_headers(self, file_base: Path): header_uri_service: str = "test_service" header_count: int = 10 headers: list[tuple[Event, bytes]] = [] + header_messages: list[Message] = [] approx_header_size: int = 0 max_file_bytes: int = 10000 @@ -151,6 +153,7 @@ def test_headers(self, file_base: Path): sequence=i, ) headers.append((event, event_payload)) + header_messages.append(message) approx_header_size += len(event.SerializeToString()) + len(event_payload) # Add some extra headers with duplicate URI's to ensure they are filtered out @@ -176,6 +179,7 @@ def test_headers(self, file_base: Path): approx_header_size, rel=0.5, ) + assert opened_writer.header_messages == header_messages # Test that headers cannot exceed the max file size with pytest.raises(RuntimeError): @@ -202,6 +206,7 @@ def test_headers(self, file_base: Path): opened_writer.write("test_path", StringValue(value=f"test_payload_{i}")) file_count = 1 + opened_writer.file_idx assert file_count > 5 # sufficient number of rollovers (params) + assert opened_writer.header_messages == header_messages for i in range(file_count): header_count_in_file: int = 0