323 lines
9.4 KiB
Python
323 lines
9.4 KiB
Python
# Copyright 2020-2023 Ternaris.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
"""Reader tests."""
|
|
|
|
# pylint: disable=redefined-outer-name
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import zstandard
|
|
|
|
from rosbags.rosbag2 import Reader, ReaderError, Writer
|
|
|
|
from .test_serde import MSG_JOINT, MSG_MAGN, MSG_MAGN_BIG, MSG_POLY
|
|
|
|
if TYPE_CHECKING:
|
|
from _pytest.fixtures import SubRequest
|
|
|
|
METADATA = """
|
|
rosbag2_bagfile_information:
|
|
version: 4
|
|
storage_identifier: sqlite3
|
|
relative_file_paths:
|
|
- db.db3{extension}
|
|
duration:
|
|
nanoseconds: 42
|
|
starting_time:
|
|
nanoseconds_since_epoch: 666
|
|
message_count: 4
|
|
topics_with_message_count:
|
|
- topic_metadata:
|
|
name: /poly
|
|
type: geometry_msgs/msg/Polygon
|
|
serialization_format: cdr
|
|
offered_qos_profiles: ""
|
|
message_count: 1
|
|
- topic_metadata:
|
|
name: /magn
|
|
type: sensor_msgs/msg/MagneticField
|
|
serialization_format: cdr
|
|
offered_qos_profiles: ""
|
|
message_count: 2
|
|
- topic_metadata:
|
|
name: /joint
|
|
type: trajectory_msgs/msg/JointTrajectory
|
|
serialization_format: cdr
|
|
offered_qos_profiles: ""
|
|
message_count: 1
|
|
compression_format: {compression_format}
|
|
compression_mode: {compression_mode}
|
|
"""
|
|
|
|
METADATA_EMPTY = """
|
|
rosbag2_bagfile_information:
|
|
version: 6
|
|
storage_identifier: sqlite3
|
|
relative_file_paths:
|
|
- db.db3
|
|
duration:
|
|
nanoseconds: 0
|
|
starting_time:
|
|
nanoseconds_since_epoch: 0
|
|
message_count: 0
|
|
topics_with_message_count: []
|
|
compression_format: ""
|
|
compression_mode: ""
|
|
files:
|
|
- duration:
|
|
nanoseconds: 0
|
|
message_count: 0
|
|
path: db.db3
|
|
starting_time:
|
|
nanoseconds_since_epoch: 0
|
|
custom_data:
|
|
key1: value1
|
|
key2: value2
|
|
"""
|
|
|
|
|
|
@pytest.fixture(params=['none', 'file', 'message'])
|
|
def bag(request: SubRequest, tmp_path: Path) -> Path:
|
|
"""Manually contruct bag."""
|
|
(tmp_path / 'metadata.yaml').write_text(
|
|
METADATA.format(
|
|
extension='' if request.param != 'file' else '.zstd',
|
|
compression_format='""' if request.param == 'none' else 'zstd',
|
|
compression_mode='""' if request.param == 'none' else request.param.upper(),
|
|
),
|
|
)
|
|
|
|
comp = zstandard.ZstdCompressor()
|
|
|
|
dbpath = tmp_path / 'db.db3'
|
|
dbh = sqlite3.connect(dbpath)
|
|
dbh.executescript(Writer.SQLITE_SCHEMA)
|
|
|
|
cur = dbh.cursor()
|
|
cur.execute(
|
|
'INSERT INTO topics VALUES(?, ?, ?, ?, ?)',
|
|
(1, '/poly', 'geometry_msgs/msg/Polygon', 'cdr', ''),
|
|
)
|
|
cur.execute(
|
|
'INSERT INTO topics VALUES(?, ?, ?, ?, ?)',
|
|
(2, '/magn', 'sensor_msgs/msg/MagneticField', 'cdr', ''),
|
|
)
|
|
cur.execute(
|
|
'INSERT INTO topics VALUES(?, ?, ?, ?, ?)',
|
|
(3, '/joint', 'trajectory_msgs/msg/JointTrajectory', 'cdr', ''),
|
|
)
|
|
cur.execute(
|
|
'INSERT INTO messages VALUES(?, ?, ?, ?)',
|
|
(1, 1, 666, MSG_POLY[0] if request.param != 'message' else comp.compress(MSG_POLY[0])),
|
|
)
|
|
cur.execute(
|
|
'INSERT INTO messages VALUES(?, ?, ?, ?)',
|
|
(2, 2, 708, MSG_MAGN[0] if request.param != 'message' else comp.compress(MSG_MAGN[0])),
|
|
)
|
|
cur.execute(
|
|
'INSERT INTO messages VALUES(?, ?, ?, ?)',
|
|
(
|
|
3,
|
|
2,
|
|
708,
|
|
MSG_MAGN_BIG[0] if request.param != 'message' else comp.compress(MSG_MAGN_BIG[0]),
|
|
),
|
|
)
|
|
cur.execute(
|
|
'INSERT INTO messages VALUES(?, ?, ?, ?)',
|
|
(4, 3, 708, MSG_JOINT[0] if request.param != 'message' else comp.compress(MSG_JOINT[0])),
|
|
)
|
|
dbh.commit()
|
|
|
|
if request.param == 'file':
|
|
with dbpath.open('rb') as ifh, (tmp_path / 'db.db3.zstd').open('wb') as ofh:
|
|
comp.copy_stream(ifh, ofh)
|
|
dbpath.unlink()
|
|
|
|
return tmp_path
|
|
|
|
|
|
def test_empty_bag(tmp_path: Path) -> None:
|
|
"""Test bags with broken fs layout."""
|
|
(tmp_path / 'metadata.yaml').write_text(METADATA_EMPTY)
|
|
dbpath = tmp_path / 'db.db3'
|
|
dbh = sqlite3.connect(dbpath)
|
|
dbh.executescript(Writer.SQLITE_SCHEMA)
|
|
|
|
with Reader(tmp_path) as reader:
|
|
assert reader.message_count == 0
|
|
assert reader.start_time == 2**63 - 1
|
|
assert reader.end_time == 0
|
|
assert reader.duration == 0
|
|
assert not list(reader.messages())
|
|
assert reader.custom_data['key1'] == 'value1'
|
|
assert reader.custom_data['key2'] == 'value2'
|
|
|
|
|
|
def test_reader(bag: Path) -> None:
|
|
"""Test reader and deserializer on simple bag."""
|
|
with Reader(bag) as reader:
|
|
assert reader.duration == 43
|
|
assert reader.start_time == 666
|
|
assert reader.end_time == 709
|
|
assert reader.message_count == 4
|
|
if reader.compression_mode:
|
|
assert reader.compression_format == 'zstd'
|
|
assert [x.id for x in reader.connections] == [1, 2, 3]
|
|
assert [*reader.topics.keys()] == ['/poly', '/magn', '/joint']
|
|
gen = reader.messages()
|
|
|
|
connection, timestamp, rawdata = next(gen)
|
|
assert connection.topic == '/poly'
|
|
assert connection.msgtype == 'geometry_msgs/msg/Polygon'
|
|
assert timestamp == 666
|
|
assert rawdata == MSG_POLY[0]
|
|
|
|
for idx in range(2):
|
|
connection, timestamp, rawdata = next(gen)
|
|
assert connection.topic == '/magn'
|
|
assert connection.msgtype == 'sensor_msgs/msg/MagneticField'
|
|
assert timestamp == 708
|
|
assert rawdata == [MSG_MAGN, MSG_MAGN_BIG][idx][0]
|
|
|
|
connection, timestamp, rawdata = next(gen)
|
|
assert connection.topic == '/joint'
|
|
assert connection.msgtype == 'trajectory_msgs/msg/JointTrajectory'
|
|
|
|
with pytest.raises(StopIteration):
|
|
next(gen)
|
|
|
|
|
|
def test_message_filters(bag: Path) -> None:
|
|
"""Test reader filters messages."""
|
|
with Reader(bag) as reader:
|
|
magn_connections = [x for x in reader.connections if x.topic == '/magn']
|
|
gen = reader.messages(connections=magn_connections)
|
|
connection, _, _ = next(gen)
|
|
assert connection.topic == '/magn'
|
|
connection, _, _ = next(gen)
|
|
assert connection.topic == '/magn'
|
|
with pytest.raises(StopIteration):
|
|
next(gen)
|
|
|
|
gen = reader.messages(start=667)
|
|
connection, _, _ = next(gen)
|
|
assert connection.topic == '/magn'
|
|
connection, _, _ = next(gen)
|
|
assert connection.topic == '/magn'
|
|
connection, _, _ = next(gen)
|
|
assert connection.topic == '/joint'
|
|
with pytest.raises(StopIteration):
|
|
next(gen)
|
|
|
|
gen = reader.messages(stop=667)
|
|
connection, _, _ = next(gen)
|
|
assert connection.topic == '/poly'
|
|
with pytest.raises(StopIteration):
|
|
next(gen)
|
|
|
|
gen = reader.messages(connections=magn_connections, stop=667)
|
|
with pytest.raises(StopIteration):
|
|
next(gen)
|
|
|
|
gen = reader.messages(start=666, stop=666)
|
|
with pytest.raises(StopIteration):
|
|
next(gen)
|
|
|
|
|
|
def test_user_errors(bag: Path) -> None:
|
|
"""Test user errors."""
|
|
reader = Reader(bag)
|
|
with pytest.raises(ReaderError, match='Rosbag is not open'):
|
|
next(reader.messages())
|
|
|
|
|
|
def test_failure_cases(tmp_path: Path) -> None:
|
|
"""Test bags with broken fs layout."""
|
|
with pytest.raises(ReaderError, match='not read metadata'):
|
|
Reader(tmp_path)
|
|
|
|
metadata = tmp_path / 'metadata.yaml'
|
|
|
|
metadata.write_text('')
|
|
with pytest.raises(ReaderError, match='not read'), \
|
|
mock.patch.object(Path, 'read_text', side_effect=PermissionError):
|
|
Reader(tmp_path)
|
|
|
|
metadata.write_text(' invalid:\nthis is not yaml')
|
|
with pytest.raises(ReaderError, match='not load YAML from'):
|
|
Reader(tmp_path)
|
|
|
|
metadata.write_text('foo:')
|
|
with pytest.raises(ReaderError, match='key is missing'):
|
|
Reader(tmp_path)
|
|
|
|
metadata.write_text(
|
|
METADATA.format(
|
|
extension='',
|
|
compression_format='""',
|
|
compression_mode='""',
|
|
).replace('version: 4', 'version: 999'),
|
|
)
|
|
with pytest.raises(ReaderError, match='version 999'):
|
|
Reader(tmp_path)
|
|
|
|
metadata.write_text(
|
|
METADATA.format(
|
|
extension='',
|
|
compression_format='""',
|
|
compression_mode='""',
|
|
).replace('sqlite3', 'hdf5'),
|
|
)
|
|
with pytest.raises(ReaderError, match='Storage plugin'):
|
|
Reader(tmp_path)
|
|
|
|
metadata.write_text(
|
|
METADATA.format(
|
|
extension='',
|
|
compression_format='""',
|
|
compression_mode='""',
|
|
),
|
|
)
|
|
with pytest.raises(ReaderError, match='files are missing'):
|
|
Reader(tmp_path)
|
|
|
|
(tmp_path / 'db.db3').write_text('')
|
|
|
|
metadata.write_text(
|
|
METADATA.format(
|
|
extension='',
|
|
compression_format='""',
|
|
compression_mode='""',
|
|
).replace('cdr', 'bson'),
|
|
)
|
|
with pytest.raises(ReaderError, match='Serialization format'):
|
|
Reader(tmp_path)
|
|
|
|
metadata.write_text(
|
|
METADATA.format(
|
|
extension='',
|
|
compression_format='"gz"',
|
|
compression_mode='"file"',
|
|
),
|
|
)
|
|
with pytest.raises(ReaderError, match='Compression format'):
|
|
Reader(tmp_path)
|
|
|
|
metadata.write_text(
|
|
METADATA.format(
|
|
extension='',
|
|
compression_format='""',
|
|
compression_mode='""',
|
|
),
|
|
)
|
|
with pytest.raises(ReaderError, match='not open database'), \
|
|
Reader(tmp_path) as reader:
|
|
next(reader.messages())
|