Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 115 additions & 0 deletions dimos/core/test_cdr_default_recording.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
# Copyright 2026 Dimensional Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""CDR passes through real mem2 storage and self-describing MCAP replay."""

from contextlib import ExitStack
from functools import partial
import threading

from dimos_cdrgen.registry import encode, schema
from mcap.reader import make_reader
from mcap.writer import Writer
import pytest
from reactivex.scheduler import ThreadPoolScheduler

from dimos.core import transport as transport_module
from dimos.core.global_config import global_config
from dimos.core.transport_factory import make_transport
from dimos.memory.store.cdr_mcap import CdrMcapStore
from dimos.memory.store.sqlite import SqliteStore
from dimos.memory.tap import TransportRecorder
from dimos.messages.builtin_interfaces.msg import Time
from dimos.messages.dimos_msgs.msg import LineSegment3D, LineSegments3D
from dimos.messages.geometry_msgs.msg import Point
from dimos.messages.std_msgs.msg import Header
from dimos.protocol.pubsub.impl.lcmpubsub import LCM
from dimos.utils.threadpool import get_max_workers


def message(sec=17):
return LineSegments3D(
header=Header(stamp=Time(sec=sec, nanosec=250_000_000), frame_id="imu"),
segments=[LineSegment3D(Point(1.0, 2.0, 3.0), Point(4.0, 5.0, 6.0), 0.75)],
)


@pytest.mark.parametrize("backend", ["lcm", "zenoh"])
def test_factory_capture_sqlite_mcap_auto_named_replay(
backend, monkeypatch, tmp_path, owned_scheduler
):
# Isolate only session configuration; codecs, pub/sub and storage stay real.
monkeypatch.setattr(transport_module, "LCM", partial(LCM, url="memq://"))
for key, value in {
"transport": backend,
"robot_ip": None,
"robot_ips": "",
"zenoh_connect": "",
"zenoh_mode": "peer",
"zenoh_multicast": False,
"zenoh_gossip": False,
"zenoh_scouting": False,
}.items():
monkeypatch.setattr(global_config, key, value)
expected = message()
expected.header.stamp.nanosec = 123456789
database = tmp_path / "live.db"
received = threading.Event()
with SqliteStore(path=str(database)) as store, ExitStack() as cleanup:
transport = make_transport("sensor/segments", LineSegments3D)
cleanup.callback(transport.stop)
recorder = TransportRecorder(store)
cleanup.callback(recorder.close)
unsubscribe = recorder.tap("segments", LineSegments3D, transport)
assert unsubscribe is not None
cleanup.callback(unsubscribe)
live = store.stream("segments", LineSegments3D).live()
cleanup.callback(live.stop)
cleanup.enter_context(live.subscribe(lambda _: received.set()))
transport.publish(expected)
assert received.wait(5), "default-factory payload was not recorded"
assert recorder.dropped == 0
path = tmp_path / "live.mcap"
with SqliteStore(path=str(database), must_exist=True) as store, path.open("wb") as raw:
writer = Writer(raw)
writer.start(profile="ros2")
captured = next(iter(store.streams.segments)).data
assert type(captured) is LineSegments3D
assert encode(captured) == encode(expected)
schema_id = writer.register_schema(
LineSegments3D.__msgtype__, "ros2msg", schema(LineSegments3D.__msgtype__).encode()
)
channel_id = writer.register_channel(transport.channel, "cdr", schema_id)
writer.add_message(channel_id, 20_000_000_000, encode(captured), 17_123456789)
writer.finish()
with CdrMcapStore(path=str(path), streams={"sensor_segments": transport.channel}) as store:
assert store.list_streams() == ["sensor_segments"]
values = list(store.replay().streams.sensor_segments.iterate_ts())
assert values == [(20.0, expected)]
assert values[0][1].header.stamp.nanosec == 123456789
with path.open("rb") as raw:
definition, channel, recorded = next(make_reader(raw).iter_messages())
assert recorded.publish_time == 17_123456789
assert recorded.log_time == 20_000_000_000
assert "MSG: std_msgs/Header" in definition.data.decode()
assert channel.topic == transport.channel


@pytest.fixture
def owned_scheduler(monkeypatch):
scheduler = ThreadPoolScheduler(max_workers=get_max_workers())
monkeypatch.setattr("dimos.utils.reactive.get_scheduler", lambda: scheduler)
monkeypatch.setattr("dimos.utils.threadpool.get_scheduler", lambda: scheduler)
yield scheduler
scheduler.executor.shutdown(wait=True)
3 changes: 1 addition & 2 deletions dimos/core/test_transport_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
from dimos.core.transport import (
DDSTransport,
LCMTransport,
ROSTransport,
ZenohTransport,
pLCMTransport,
pZenohTransport,
Expand Down Expand Up @@ -177,7 +176,7 @@ def test_factory_preserves_pickle_for_python_types(config, message_type):
assert isinstance(transport, expected)


@pytest.mark.parametrize("transport_type", [ROSTransport, DDSTransport])
@pytest.mark.parametrize("transport_type", [DDSTransport])
def test_generated_messages_require_ros_dds_type_adapters(transport_type):
with pytest.raises(TypeError, match="type adapter"):
transport_type("generated", Int32)
Expand Down
59 changes: 59 additions & 0 deletions dimos/imitation/collection/test_cdr_default_graph.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
# Copyright 2026 Dimensional Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest

from dimos.core.coordination.blueprints import Blueprint
from dimos.imitation.collection.blueprint import (
learning_collect_webxr_piper,
learning_collect_webxr_xarm7,
)
from dimos.messages.sensor_msgs.msg import JointState
from dimos.messages.std_msgs.msg import String
from dimos.teleop.webxr.extensions import ArmTeleopModule

AGGREGATE = "coordinator_joint_state"


@pytest.mark.parametrize(
"blueprint",
[learning_collect_webxr_xarm7, learning_collect_webxr_piper],
)
def test_collection_status_is_wired_to_webxr_hud(blueprint: Blueprint) -> None:
hud = next(atom for atom in blueprint.blueprints if atom.module is ArmTeleopModule)
status = next(stream for stream in hud.streams if stream.name == "status")

assert status.direction == "in"
assert status.type is String


def _joint_streams(blueprint: Blueprint) -> dict[tuple[str, str], str]:
"""(instance, port) -> effective stream name, as ModuleCoordinator pairs streams."""
return {
(atom.name, stream.name): blueprint.remapping_map.get((atom.name, stream.name), stream.name)
for atom in blueprint.active_blueprints
for stream in atom.streams
if stream.type is JointState
}


@pytest.mark.parametrize("blueprint", [learning_collect_webxr_xarm7, learning_collect_webxr_piper])
def test_recorder_reads_aggregate_joint_state(blueprint: Blueprint) -> None:
streams = _joint_streams(blueprint)

# Plain name pairing on both ends, no remap in between. The coordinator
# atom carries its explicit instance_name (the RPC lookup contract).
assert streams[("collectionrecorder", AGGREGATE)] == AGGREGATE
assert streams[("ControlCoordinator", AGGREGATE)] == AGGREGATE
assert not [port for _instance, port in streams if port.endswith("_joints")]
35 changes: 35 additions & 0 deletions dimos/robot/galaxea/r1pro/test_cdr_default_graph.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
# Copyright 2025-2026 Dimensional Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from contextlib import ExitStack

from dimos.core.global_config import global_config
from dimos.core.transport_factory import make_transport
from dimos.messages.geometry_msgs.msg import PoseStamped, Twist
from dimos.robot.galaxea.r1pro.blueprints.basic.r1pro_coordinator import r1pro_control


def test_r1_blueprint_chassis_pins_match_default_factory(monkeypatch):
blueprint = r1pro_control()
# The existing R1 control subtree selects Zenoh for both adapter and driver.
assert blueprint.global_config_overrides["transport"] == "zenoh"
monkeypatch.setattr(global_config, "transport", "zenoh")
with ExitStack() as cleanup:
for name, kind in [("cmd_vel", Twist), ("odom", PoseStamped)]:
pin = blueprint.transport_map[("chassis_" + name, kind)].build()
default = make_transport("/chassis/" + name, kind)
cleanup.callback(pin.stop)
cleanup.callback(default.stop)
assert pin.channel == default.channel
assert type(pin) is type(default)
Loading