From 4dd42a3aa10d19fa5c9cd8cbd2a353926c6ecfa7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Wed, 3 Jun 2026 22:12:29 -0300 Subject: [PATCH 01/11] Support runtime add/remove of robots for fleet autodiscovery Make `FleetConnector.update_fleet()` the single, idempotent reconciler of the live robot fleet, and add `add_robot()`/`remove_robot()` as thin convenience wrappers over it. This lets fleet connectors add and remove robots while running, enabling robot autodiscovery (e.g. polling a fleet manager and reacting to robots joining/leaving). Previously the fleet was static after `start()`: `update_fleet()` only swapped the cached config and `__robot_ids`, with sessions bulk-created once at connect time and never reconciled afterwards. Changes: - `update_fleet(fleet)` now reconciles the live sessions to exactly the given list: creates+connects sessions for added robots (registering cameras, command handler and online-status callback) and disconnects + frees sessions for removed robots, clearing their per-robot state. The whole reconcile runs under a re-entrant fleet lock so a fleet change is atomic w.r.t. concurrent callers. - `add_robot()`/`remove_robot()` read the current fleet and delegate to `update_fleet()` under the same lock. - Reconcile logic lives in a private `__apply_fleet()`; `update_fleet()` wraps it and `__connect()` calls it directly to create the initial sessions (replacing `__initialize_sessions()`). - Camera registration moved from `__run_connector` into `__initialize_session` so runtime-added robots get their cameras too. - `__disconnect()` disconnects under the lock and clears the session dict (so a later `start()` reconnects, and a concurrent add can't leak a session mid-teardown). - `__publish_pending_system_stats` reads `__robot_sessions` instead of the pool, so a removed robot is skipped rather than resurrected on a cache miss. - Single-robot `Connector` overrides the public mutation API to raise `NotImplementedError` (its startup still works via the private `__apply_fleet`). - Fleet may shrink to zero robots at runtime; the loop idles safely. Docs and the fleet-connector example updated to show autodiscovery. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/contents/usage/fleet.md | 52 ++++- examples/fleet-connector/connector.py | 16 +- inorbit_connector/connector.py | 271 +++++++++++++++++------ tests/test_connector.py | 299 +++++++++++++++++++++++++- 4 files changed, 559 insertions(+), 79 deletions(-) diff --git a/docs/contents/usage/fleet.md b/docs/contents/usage/fleet.md index 00e783d..10a9ccf 100644 --- a/docs/contents/usage/fleet.md +++ b/docs/contents/usage/fleet.md @@ -69,24 +69,56 @@ for robot_id in self.robot_ids: ### Updating the Fleet -You can update the fleet configuration during `_connect()` by calling the `update_fleet()` method. This is useful for dynamically setting the robots list before the connector starts, for example, when provisioning robots from fleet manager data instead of hardcoded values in the config files: +`update_fleet()` reconciles the connector's live fleet to exactly the list you pass: it creates and connects a session for each newly added robot (registering its cameras, command handler and online-status callback) and disconnects and frees the session of each robot that is no longer present. It is idempotent — robots already in the fleet are left untouched. + +This makes it the entry point both for **initial provisioning** (declaring the fleet from fleet-manager data instead of hardcoded config) and for **runtime autodiscovery**: ```python @override async def _connect(self) -> None: - """Connect to fleet manager and update fleet.""" + """Connect to fleet manager and declare the fleet.""" # Fetch robot list from fleet manager API robots = await self._fleet_manager.get_robots() - - # Update fleet configuration - fleet_config = [ - RobotConfig(robot_id=robot.id, cameras=robot.cameras) - for robot in robots - ] - self.update_fleet(fleet_config) + + # Reconcile the fleet to the discovered robots + self.update_fleet( + [RobotConfig(robot_id=robot.id, cameras=robot.cameras) for robot in robots] + ) +``` + +#### Adding and removing robots at runtime + +For targeted, event-driven changes (e.g. a robot joining or leaving the fleet while the connector runs), use `add_robot()` and `remove_robot()`: + +```python +# A robot appeared — create and connect its session immediately +self.add_robot(RobotConfig(robot_id="robot-42", cameras=[])) + +# A robot left — disconnect and free its session, clearing its state +self.remove_robot("robot-42") +``` + +A typical autodiscovery loop diffs the fleet manager's current robot set against `self.robot_ids`: + +```python +@override +async def _execution_loop(self) -> None: + discovered = set(await self._fleet_manager.fetch_robot_list()) + current = set(self.robot_ids) + for robot_id in discovered - current: + self.add_robot(RobotConfig(robot_id=robot_id)) + for robot_id in current - discovered: + self.remove_robot(robot_id) + ... # publish data for self.robot_ids ``` -The `update_fleet()` method updates the fleet configuration and initializes sessions for all robots. +:::{note} +- These methods create or destroy sessions **immediately**, so call them once the connector is connecting or running (from `_connect()` onward), not before `start()`. +- `add_robot()` raises `ValueError` on a duplicate `robot_id`; `remove_robot()` is a no-op (logs a warning) for an unknown id, so it is safe to call from a discovery loop that may fire repeatedly. +- The fleet may shrink to zero robots at runtime; the execution loop simply idles until robots are discovered again. (The loaded configuration still requires at least one robot at startup.) +- A robot present in both the old and new fleet is treated as unchanged even if its `RobotConfig` differs (e.g. its cameras changed). To apply a changed config to a running robot, call `remove_robot()` then `add_robot()`. +- These methods are thread-safe and may be called from the execution loop, a command handler, or any other thread. +::: ## Publishing Methods diff --git a/examples/fleet-connector/connector.py b/examples/fleet-connector/connector.py index 96637c8..3b4483e 100644 --- a/examples/fleet-connector/connector.py +++ b/examples/fleet-connector/connector.py @@ -70,10 +70,18 @@ async def _connect(self) -> None: """ self._fleet_manager.start() - # Here the connector may fetch a robot list from the fleet manager. e.g.: - # robots: list[RobotConfig] = self._fleet_manager.fetch_robot_list() - # self.update_fleet(robots) - # `robots` will be added to the InOrbit fleet + # The connector may fetch the robot list from the fleet manager here and + # declare it as the InOrbit fleet via update_fleet(). update_fleet() sets + # the live sessions to the given list, so it can be called both here for + # initial provisioning and later at runtime in cases of fleet autodiscovery + # e.g.: + # + # robots = await self._fleet_manager.fetch_robot_list() + # self.update_fleet( + # [RobotConfig(robot_id=r.id, cameras=r.cameras) for r in robots] + # ) + # + # For runtime changes use add_robot()/remove_robot() instead. self._logger.info( f"Connected to fleet manager API {self.api_version} for " diff --git a/inorbit_connector/connector.py b/inorbit_connector/connector.py index a006fb9..2a4cfc2 100644 --- a/inorbit_connector/connector.py +++ b/inorbit_connector/connector.py @@ -99,13 +99,12 @@ def __init__(self, config: ConnectorRootConfig, **kwargs) -> None: # Common information self.config = config - # Cache of robot IDs in config.fleet. Accessed through the robot_ids property - # Updated by update_fleet() - self.__robot_ids: list[str] = [] - # update_fleet() may be called during user-defined implementation of _connect() - # to update the fleet before initializing the robot sessions - # Initialize the robot_ids cache - self.update_fleet(config.fleet) + # Cache of robot IDs in config.fleet. Accessed through the robot_ids property. + # Seeded directly from the config here (no sessions are created during + # construction); thereafter kept in sync by update_fleet()/add_robot()/ + # remove_robot(). The initial robot sessions are created later by __connect(), + # which calls update_fleet(self.config.fleet) once the connector thread is up. + self.__robot_ids: list[str] = [robot.robot_id for robot in config.fleet] # Per robot state self.__last_published_frame_ids: dict[str, str] = {} @@ -117,10 +116,13 @@ def __init__(self, config: ConnectorRootConfig, **kwargs) -> None: # Managed temp directory for fetched map files (lazily initialized) self.__temp_map_dir: tempfile.TemporaryDirectory | None = None - # Private dictionary for fast internal access (use self._get_session(robot_id) - # for thread-safe access) in tight loops. It should not be accessed directly - # by subclasses to maintain thread-safety + # Private dictionary for fast internal access (use self._get_robot_session( + # robot_id) for thread-safe access) in tight loops. It should not be accessed + # directly by subclasses to maintain thread-safety self.__robot_sessions: dict[str, RobotSession] = {} + # Lock guarding fleet mutation: __robot_ids, __robot_sessions membership, + # config.fleet, and the per-robot state dicts. + self.__fleet_lock = threading.RLock() # Threading for the main run methods # The connector runs an asycio loop within a spawned thread @@ -211,8 +213,10 @@ def __init__(self, config: ConnectorRootConfig, **kwargs) -> None: @property def robot_ids(self) -> list[str]: """Get the list of robot IDs in the fleet.""" - # Return the cached list of robot IDs - return self.__robot_ids + # Return a copy of the cached list so callers can't mutate internal state and + # so a captured reference stays consistent if the fleet changes concurrently. + with self.__fleet_lock: + return self.__robot_ids.copy() @property def _connector_type(self) -> str: @@ -228,17 +232,133 @@ def _connector_type(self) -> str: def update_fleet(self, fleet: list[RobotConfig]) -> None: """Update the robot fleet. - This method may be called during the user-defined implementation of _connect() - to update the fleet configuration before initializing the robot sessions. - e.g. fetching the robot list from a fleet manager API. + This is the single entry point for fleet membership. It diffs the requested + fleet against the currently active robot sessions and: + + - creates and connects a session for each newly added robot (registering its + cameras, command handler and online-status callback), and + - disconnects and frees the session of each robot no longer present, clearing + its per-robot state. + + It is idempotent: robots already in the fleet are left untouched. + ``__connect()`` calls it with ``self.config.fleet`` to create the + initial sessions, and subclasses may call it at runtime (from ``_connect()`` + onward) to implement robot autodiscovery. e.g. fetching the robot list from a + fleet manager API. + + Note: + Sessions are created/destroyed immediately, so this must be called once the + connector is connecting or running (i.e. from ``_connect()`` or the + execution loop), not before ``start()``. + + A robot present in both the old and new fleet is treated as unchanged even + if its ``RobotConfig`` differs (e.g. its cameras changed). To apply a + changed config to a running robot, call ``remove_robot()`` then + ``add_robot()``. Args: fleet (list[RobotConfig]): The new fleet configuration + + Raises: + ValueError: If ``fleet`` contains duplicate robot IDs. + """ + self.__apply_fleet(fleet) + + def __apply_fleet(self, fleet: list[RobotConfig]) -> None: + """Reconcile the live robot sessions to match ``fleet``. + + This is the implementation behind ``update_fleet()``. It is kept private so the + connector's own startup (``__connect``) can reconcile the initial fleet even on + the single-robot ``Connector`` subclass, which overrides the public + ``update_fleet``/``add_robot``/``remove_robot`` to raise. """ - # Update the fleet configuration - self.config.fleet = fleet - # Update robot ID cache - self.__robot_ids = [robot.robot_id for robot in self.config.fleet] + new_ids = [robot.robot_id for robot in fleet] + if len(set(new_ids)) != len(new_ids): + raise ValueError("Robot ids must be unique") + + # The whole reconcile runs under the lock so a fleet change is atomic with + # respect to concurrent add_robot()/remove_robot()/update_fleet() calls. The + # lock is re-entrant, so the wrappers can hold it across their read+delegate. + with self.__fleet_lock: + new_id_set = set(new_ids) + to_remove = [rid for rid in self.__robot_sessions if rid not in new_id_set] + to_add = [rid for rid in new_ids if rid not in self.__robot_sessions] + + # Apply the new membership before touching sessions so the running loop + # sees the target fleet and stops iterating removed robots. + self.config.fleet = list(fleet) + self.__robot_ids = list(new_ids) + + for robot_id in to_remove: + self.__last_published_frame_ids.pop(robot_id, None) + self.__pending_system_stats.pop(robot_id, None) + self.__robot_sessions.pop(robot_id, None) + # free_robot_session disconnects and removes the session from the + # pool; no-op if the pool has no session for this robot. + self.__session_pool.free_robot_session(robot_id) + for robot_id in to_add: + self.__robot_sessions[robot_id] = self.__initialize_session(robot_id) + + def add_robot(self, robot_config: RobotConfig) -> None: + """Add a single robot to the fleet at runtime. + + Convenience wrapper over ``update_fleet()``: appends the robot to the current + fleet and reconciles, which creates and connects its session (registering its + cameras, command handler and online-status callback). + + Note: + The session is created immediately, so call this once the connector is + connecting or running (from ``_connect()`` onward), not before ``start()``. + + Args: + robot_config (RobotConfig): The configuration of the robot to add. + + Raises: + ValueError: If a robot with the same ``robot_id`` is already in the fleet. + """ + with self.__fleet_lock: + if robot_config.robot_id in self.__robot_ids: + raise ValueError( + f"Robot '{robot_config.robot_id}' is already in the fleet" + ) + self.update_fleet([*self.config.fleet, robot_config]) + + def remove_robot(self, robot_id: str) -> None: + """Remove a single robot from the fleet at runtime. + + Convenience wrapper over ``update_fleet()``: drops the robot from the current + fleet and reconciles, which disconnects and frees its session and clears its + per-robot state. Idempotent: removing a robot that is not in the fleet logs a + warning and returns without error. + + Args: + robot_id (str): The ID of the robot to remove. + """ + with self.__fleet_lock: + if robot_id not in self.__robot_ids: + self._logger.warning( + f"remove_robot: robot '{robot_id}' is not in the fleet" + ) + return + self.update_fleet( + [robot for robot in self.config.fleet if robot.robot_id != robot_id] + ) + + def _get_robot_config(self, robot_id: str) -> RobotConfig | None: + """Return the ``RobotConfig`` for ``robot_id`` from the fleet, or None. + + Args: + robot_id (str): The robot ID to look up. + + Returns: + RobotConfig | None: The matching configuration, or None if the robot is not + in the fleet. + """ + with self.__fleet_lock: + for robot_config in self.config.fleet: + if robot_config.robot_id == robot_id: + return robot_config + return None def _handle_command_exception( self, @@ -356,24 +476,30 @@ def __initialize_session(self, robot_id: str) -> RobotSession: session, self._inorbit_robot_command_handler ) - return session - - def __initialize_sessions(self) -> None: - """Initialize the robot sessions.""" + # Register cameras declared in this robot's configuration. + robot_config = self._get_robot_config(robot_id) + if robot_config is not None: + for idx, camera_config in enumerate(robot_config.cameras): + self._logger.info( + f"Registering camera {idx} for robot {robot_id}: " + f"{str(camera_config.video_url)}" + ) + # If values are None, remove the key from the dictionary to use + # edge-sdk defaults + dump = camera_config.model_dump() + clean = {k: v for k, v in dump.items() if v is not None} + session.register_camera(str(idx), OpenCVCamera(**clean)) - for robot_id in self.robot_ids: - self.__robot_sessions[robot_id] = self.__initialize_session(robot_id) - self._logger.info( - f"Initialized {len(self.__robot_sessions)} robot sessions for robots " - f"{', '.join(self.robot_ids)}" - ) + return session async def __connect(self) -> None: """Initialize the connection to InOrbit based on the provided configuration, and connect to external services calling self._connect(). - self.update_fleet() may be called during this method to update the fleet - configuration before initializing the robot sessions. + After self._connect() returns, the initial robot sessions are created via + self.update_fleet(self.config.fleet). Subclasses may also populate the fleet + themselves during self._connect() (via update_fleet() or add_robot()), in which + case that final reconcile is a no-op. Raises: Exception: If the robot session cannot connect. @@ -381,15 +507,21 @@ async def __connect(self) -> None: # Call the user-implemented connection logic await self._connect() - # Connect to InOrbit - self.__initialize_sessions() + # Create the InOrbit sessions for the current fleet. Calls the private reconcile + # directly so it also works on the single-robot Connector subclass, which blocks + # the public update_fleet(). + self.__apply_fleet(self.config.fleet) async def __disconnect(self) -> None: """Disconnect external services and disconnect from InOrbit.""" - # Disconnect from InOrbit - for session in self.__robot_sessions.values(): - session.disconnect() + # Disconnect from InOrbit while holding the lock so no robot can be added (via + # add_robot/update_fleet from another thread) while we tear the fleet down, and + # clear the session dict so a subsequent start() recreates the sessions. + with self.__fleet_lock: + for session in self.__robot_sessions.values(): + session.disconnect() + self.__robot_sessions.clear() # Clean up temporary map files if self.__temp_map_dir is not None: @@ -410,25 +542,9 @@ def __run_connector(self) -> None: self.__loop = asyncio.new_event_loop() asyncio.set_event_loop(self.__loop) - # Connect to external services and create the InOrbit session - self.__loop.run_until_complete(self.__connect()) - - # Set up camera feeds - for robot_config in self.config.fleet: - for idx, camera_config in enumerate(robot_config.cameras): - self._logger.info( - f"Registering camera {idx} for robot {robot_config.robot_id}: " - f"{str(camera_config.video_url)}" - ) - # If values are None, remove the key from the dictionary to use - # edge-sdk defaults - dump = camera_config.model_dump() - clean = {k: v for k, v in dump.items() if v is not None} - self.__robot_sessions[robot_config.robot_id].register_camera( - str(idx), OpenCVCamera(**clean) - ) - try: + # Connect to external services and create the InOrbit sessions. + self.__loop.run_until_complete(self.__connect()) self.__loop.run_until_complete(self.__run_loop()) except Exception as e: self._logger.error(f"Error in execution loop: {e}") @@ -767,7 +883,8 @@ def publish_robot_system_stats(self, robot_id: str, **kwargs) -> None: **kwargs: System stats data (cpu_load_percentage, ram_usage_percentage, hdd_usage_percentage, ts) """ - self.__pending_system_stats[robot_id] = kwargs + with self.__fleet_lock: + self.__pending_system_stats[robot_id] = kwargs def __get_connector_system_stats(self) -> dict: """Get system stats from the connector's host environment. @@ -807,16 +924,21 @@ def __publish_pending_system_stats(self) -> None: } ) - for robot_id in self.robot_ids: - session = self._get_robot_session(robot_id) - if pending_status := self.__pending_system_stats.get(robot_id): + with self.__fleet_lock: + sessions = { + robot_id: self.__robot_sessions[robot_id] + for robot_id in self.__robot_ids + if robot_id in self.__robot_sessions + } + pending = self.__pending_system_stats + self.__pending_system_stats = {} + + for robot_id, session in sessions.items(): + if pending_status := pending.get(robot_id): session.publish_system_stats(**pending_status) else: session.publish_system_stats(**default_values) - # Clear stored stats for next iteration - self.__pending_system_stats.clear() - # Methods meant to be extended by subclasses @abstractmethod async def _connect(self) -> None: @@ -825,8 +947,10 @@ async def _connect(self) -> None: This method should not be called directly. Instead, call the start() method to start the connector. This ensures that the connector is only started once. - self.update_fleet() may be called during this method to update the fleet - configuration before initializing the robot sessions. + self.update_fleet() (or add_robot()) may be called during this method to declare + the fleet from an external source, e.g. fetching the robot list from a fleet + manager API. Doing so creates the robot sessions immediately; otherwise they are + created right after this method returns using data from the fleet configuration. """ ... @@ -913,6 +1037,27 @@ class Connector(FleetConnector, ABC): FleetConnector.__init__() for more details. """ + @override + def update_fleet(self, fleet: list[RobotConfig]) -> None: + raise NotImplementedError( + "udpate_fleet() is not supported on a single-robot Connector; " + "use FleetConnector for multi-robot fleets" + ) + + @override + def add_robot(self, robot_config: RobotConfig) -> None: + raise NotImplementedError( + "add_robot() is not supported on a single-robot Connector; " + "use FleetConnector for multi-robot fleets" + ) + + @override + def remove_robot(self, robot_id: str) -> None: + raise NotImplementedError( + "remove_robot() is not supported on a single-robot Connector; " + "use FleetConnector for multi-robot fleets" + ) + def __init__(self, robot_id: str, config: ConnectorRootConfig, **kwargs) -> None: """Initialize a new InOrbit connector. diff --git a/tests/test_connector.py b/tests/test_connector.py index 94e8f15..d85e7bf 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -7,13 +7,16 @@ # Standard import os +import threading from time import sleep from unittest.mock import AsyncMock, MagicMock, patch # Third-party import pytest from pydantic import AnyHttpUrl +from inorbit_edge.models import CameraConfig from inorbit_edge.robot import RobotSession +from inorbit_edge.video import OpenCVCamera # InOrbit from inorbit_connector.connector import ( @@ -591,8 +594,8 @@ def make_fleet_connector_not_abstract(self): FleetConnector.__abstractmethods__ = set() @pytest.fixture - def fleet_connector(self, base_model): - return FleetConnector( + def fleet_connector(self, base_model, mock_robot_session_pool): + connector = FleetConnector( ConnectorRootConfig( **base_model, fleet=[ @@ -601,6 +604,13 @@ def fleet_connector(self, base_model): ], ) ) + # Simulate the post-connect state: sessions are created by update_fleet() + # during __connect(). __publish_pending_system_stats reads from this dict. + connector._FleetConnector__robot_sessions = { + robot_id: mock_robot_session_pool.get_session(robot_id) + for robot_id in connector.robot_ids + } + return connector def test_publish_pending_system_stats_publishes_stored_stats( self, fleet_connector, mock_robot_session_pool @@ -709,6 +719,10 @@ def test_publish_connector_system_stats_uses_psutil( ), publish_connector_system_stats=True, ) + # Simulate the post-connect state (sessions created by update_fleet). + connector._FleetConnector__robot_sessions = { + "TestRobot1": mock_robot_session_pool.get_session("TestRobot1") + } connector._FleetConnector__publish_pending_system_stats() @@ -731,6 +745,10 @@ def test_publish_connector_system_stats_fallback_without_psutil( ), publish_connector_system_stats=True, ) + # Simulate the post-connect state (sessions created by update_fleet). + connector._FleetConnector__robot_sessions = { + "TestRobot1": mock_robot_session_pool.get_session("TestRobot1") + } # Should have fallen back to False assert connector._FleetConnector__publish_connector_system_stats is False @@ -775,6 +793,254 @@ def test_publish_connector_system_stats_disabled_by_default( assert fleet_connector._FleetConnector__publish_connector_system_stats is False +class TestFleetConnectorRuntimeFleet: + """Tests for runtime fleet add/remove (autodiscovery support).""" + + @pytest.fixture + def base_model(self): + return { + "api_key": "valid_key", + "connection_config_url": AnyHttpUrl("https://valid.com/"), + "connector_type": "valid_connector", + "connector_config": DummyConfig(), + } + + @pytest.fixture(autouse=True) + def make_fleet_connector_not_abstract(self): + FleetConnector.__abstractmethods__ = set() + + @pytest.fixture + def fleet_connector(self, base_model, mock_robot_session_pool): + """A connector in the post-connect state (sessions created for the fleet).""" + connector = FleetConnector( + ConnectorRootConfig( + **base_model, + fleet=[ + RobotConfig(robot_id="TestRobot1"), + RobotConfig(robot_id="TestRobot2"), + ], + ) + ) + # Simulate the state after __connect() -> update_fleet() created the sessions. + connector._FleetConnector__robot_sessions = { + robot_id: mock_robot_session_pool.get_session(robot_id) + for robot_id in connector.robot_ids + } + return connector + + @staticmethod + def _sessions(connector): + return connector._FleetConnector__robot_sessions + + def test_add_robot_creates_and_connects_session( + self, fleet_connector, mock_robot_session_pool + ): + fleet_connector.add_robot(RobotConfig(robot_id="R3")) + + assert "R3" in fleet_connector.robot_ids + assert any(rc.robot_id == "R3" for rc in fleet_connector.config.fleet) + assert "R3" in self._sessions(fleet_connector) + mock_robot_session_pool.get_session.assert_any_call("R3", robot_name="R3") + + def test_add_robot_duplicate_raises( + self, fleet_connector, mock_robot_session_pool + ): + mock_robot_session_pool.get_session.reset_mock() + with pytest.raises(ValueError): + fleet_connector.add_robot(RobotConfig(robot_id="TestRobot1")) + + # Fleet unchanged and no new session created. + assert fleet_connector.robot_ids == ["TestRobot1", "TestRobot2"] + mock_robot_session_pool.get_session.assert_not_called() + + def test_remove_robot_frees_session_and_clears_state( + self, fleet_connector, mock_robot_session_pool + ): + # Seed per-robot state that must be cleaned up. + fleet_connector._FleetConnector__last_published_frame_ids["TestRobot1"] = "map" + fleet_connector._FleetConnector__pending_system_stats["TestRobot1"] = {"x": 1} + + fleet_connector.remove_robot("TestRobot1") + + mock_robot_session_pool.free_robot_session.assert_called_once_with("TestRobot1") + assert "TestRobot1" not in fleet_connector.robot_ids + assert all(rc.robot_id != "TestRobot1" for rc in fleet_connector.config.fleet) + assert "TestRobot1" not in self._sessions(fleet_connector) + assert ( + "TestRobot1" + not in fleet_connector._FleetConnector__last_published_frame_ids + ) + assert ( + "TestRobot1" not in fleet_connector._FleetConnector__pending_system_stats + ) + + def test_remove_robot_unknown_is_noop( + self, fleet_connector, mock_robot_session_pool, caplog + ): + fleet_connector.remove_robot("does-not-exist") + + mock_robot_session_pool.free_robot_session.assert_not_called() + assert fleet_connector.robot_ids == ["TestRobot1", "TestRobot2"] + assert "not in the fleet" in caplog.text + + def test_update_fleet_reconciles_diff( + self, fleet_connector, mock_robot_session_pool + ): + mock_robot_session_pool.get_session.reset_mock() + mock_robot_session_pool.free_robot_session.reset_mock() + + fleet_connector.update_fleet( + [RobotConfig(robot_id="TestRobot2"), RobotConfig(robot_id="R3")] + ) + + # TestRobot1 dropped, R3 added, ordering preserved. + assert fleet_connector.robot_ids == ["TestRobot2", "R3"] + mock_robot_session_pool.free_robot_session.assert_called_once_with("TestRobot1") + mock_robot_session_pool.get_session.assert_any_call("R3", robot_name="R3") + # TestRobot2 untouched: not re-created, not freed. + assert "TestRobot2" in self._sessions(fleet_connector) + assert ("TestRobot2",) not in [ + c.args for c in mock_robot_session_pool.free_robot_session.call_args_list + ] + + def test_update_fleet_duplicate_ids_raises(self, fleet_connector): + with pytest.raises(ValueError): + fleet_connector.update_fleet( + [RobotConfig(robot_id="R3"), RobotConfig(robot_id="R3")] + ) + + def test_connect_materializes_initial_sessions( + self, base_model, mock_robot_session_pool + ): + """__connect() creates a session per initial robot via update_fleet.""" + connector = FleetConnector( + ConnectorRootConfig( + **base_model, + fleet=[ + RobotConfig(robot_id="TestRobot1"), + RobotConfig(robot_id="TestRobot2"), + ], + ) + ) + connector.start() + try: + sleep(0.5) + mock_robot_session_pool.get_session.assert_any_call( + "TestRobot1", robot_name="TestRobot1" + ) + mock_robot_session_pool.get_session.assert_any_call( + "TestRobot2", robot_name="TestRobot2" + ) + mock_robot_session_pool.free_robot_session.assert_not_called() + finally: + connector.stop() + + def test_disconnect_clears_sessions_for_restart( + self, base_model, mock_robot_session_pool + ): + """stop() disconnects + clears sessions so a later start() reconnects.""" + connector = FleetConnector( + ConnectorRootConfig( + **base_model, + fleet=[RobotConfig(robot_id="TestRobot1")], + ) + ) + connector.start() + sleep(0.5) + connector.stop() + + # Session was created and then disconnected, and the dict cleared. + session = mock_robot_session_pool.get_session("TestRobot1") + session.disconnect.assert_called() + assert connector._FleetConnector__robot_sessions == {} + + # Restart must recreate the session (to_add keys off __robot_sessions). + session.reset_mock() + connector.start() + try: + sleep(0.5) + assert "TestRobot1" in connector._FleetConnector__robot_sessions + finally: + connector.stop() + + def test_runtime_empty_fleet_does_not_crash( + self, fleet_connector, mock_robot_session_pool + ): + fleet_connector.remove_robot("TestRobot1") + fleet_connector.remove_robot("TestRobot2") + + assert fleet_connector.robot_ids == [] + assert self._sessions(fleet_connector) == {} + # Publishing pending stats over an empty fleet must be a safe no-op. + fleet_connector._FleetConnector__publish_pending_system_stats() + + def test_cameras_registered_on_runtime_add( + self, fleet_connector, mock_robot_session_pool + ): + robot = RobotConfig( + robot_id="CamBot", + cameras=[CameraConfig(video_url="rtsp://example/stream")], + ) + fleet_connector.add_robot(robot) + + session = mock_robot_session_pool.get_session("CamBot") + assert session.register_camera.call_count == 1 + args = session.register_camera.call_args + assert args.args[0] == "0" + assert isinstance(args.args[1], OpenCVCamera) + + def test_cameras_registered_once_at_startup( + self, base_model, mock_robot_session_pool + ): + """Regression: moving camera reg into __initialize_session keeps it once.""" + connector = FleetConnector( + ConnectorRootConfig( + **base_model, + fleet=[ + RobotConfig( + robot_id="CamBot", + cameras=[CameraConfig(video_url="rtsp://example/stream")], + ) + ], + ) + ) + connector.start() + try: + sleep(0.5) + session = mock_robot_session_pool.get_session("CamBot") + assert session.register_camera.call_count == 1 + finally: + connector.stop() + + def test_get_robot_config(self, fleet_connector): + assert fleet_connector._get_robot_config("TestRobot1").robot_id == "TestRobot1" + assert fleet_connector._get_robot_config("nope") is None + + def test_thread_safety_smoke(self, fleet_connector): + """Concurrent add/remove on disjoint ids leaves consistent state.""" + + def worker(start): + for i in range(start, start + 20): + rid = f"T{i}" + fleet_connector.add_robot(RobotConfig(robot_id=rid)) + fleet_connector.remove_robot(rid) + + threads = [ + threading.Thread(target=worker, args=(base,)) + for base in (100, 200, 300, 400) + ] + for t in threads: + t.start() + for t in threads: + t.join() + + # All transient robots removed; the original fleet remains; no duplicates. + ids = fleet_connector.robot_ids + assert sorted(ids) == ["TestRobot1", "TestRobot2"] + assert len(ids) == len(set(ids)) + assert set(self._sessions(fleet_connector)) == {"TestRobot1", "TestRobot2"} + + # ============================================================================== # Connector Tests (Single Robot - Subclass of FleetConnector) # ============================================================================== @@ -888,6 +1154,35 @@ def test_get_session(self, base_connector, mock_robot_session_pool): session = base_connector._get_session() assert session is not None assert session.robot_id == "TestRobot" + + def test_fleet_mutation_methods_raise(self, base_connector): + """Single-robot Connector blocks the runtime fleet-mutation API.""" + with pytest.raises(NotImplementedError): + base_connector.update_fleet([RobotConfig(robot_id="X")]) + with pytest.raises(NotImplementedError): + base_connector.add_robot(RobotConfig(robot_id="X")) + with pytest.raises(NotImplementedError): + base_connector.remove_robot("TestRobot") + + def test_start_still_creates_session_despite_blocked_update_fleet( + self, base_model, mock_robot_session_pool + ): + """__connect uses the private reconcile, so startup works though + update_fleet is blocked.""" + connector = Connector( + "TestRobot", + ConnectorRootConfig( + **base_model, fleet=[RobotConfig(robot_id="TestRobot")] + ), + ) + connector.start() + try: + sleep(0.5) + mock_robot_session_pool.get_session.assert_any_call( + "TestRobot", robot_name="TestRobot" + ) + finally: + connector.stop() mock_robot_session_pool.get_session.assert_called() def test_publish_map(self, base_model, mock_robot_session_pool): From b6cdd032dffb1f90b51c424292eb9be6cf05f009 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Wed, 3 Jun 2026 22:15:44 -0300 Subject: [PATCH 02/11] Update fleet.md --- docs/contents/usage/fleet.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/contents/usage/fleet.md b/docs/contents/usage/fleet.md index 10a9ccf..0c7ae60 100644 --- a/docs/contents/usage/fleet.md +++ b/docs/contents/usage/fleet.md @@ -69,7 +69,7 @@ for robot_id in self.robot_ids: ### Updating the Fleet -`update_fleet()` reconciles the connector's live fleet to exactly the list you pass: it creates and connects a session for each newly added robot (registering its cameras, command handler and online-status callback) and disconnects and frees the session of each robot that is no longer present. It is idempotent — robots already in the fleet are left untouched. +`update_fleet()` reconciles the connector's live fleet to exactly the list you passed. It creates and connects a session for each newly added robot, and disconnects and frees the session of each robot that is no longer present. Robots already in the fleet are left untouched. This makes it the entry point both for **initial provisioning** (declaring the fleet from fleet-manager data instead of hardcoded config) and for **runtime autodiscovery**: @@ -210,4 +210,3 @@ The lifecycle methods are the same as single-robot connectors: - **Simple fleet connector**: [examples/simple-fleet-connector/connector.py](https://github.com/inorbit-ai/inorbit-connector-python/blob/main/examples/simple-fleet-connector/connector.py) - **Fleet connector (CLI)**: [examples/fleet-connector/](https://github.com/inorbit-ai/inorbit-connector-python/tree/main/examples/fleet-connector) - **Examples index**: [examples/README.md](https://github.com/inorbit-ai/inorbit-connector-python/blob/main/examples/README.md) - From cf7573569ec7d851355f215d8acd4905e6337953 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Wed, 3 Jun 2026 22:17:55 -0300 Subject: [PATCH 03/11] Update fleet.md --- docs/contents/usage/fleet.md | 22 ++++------------------ 1 file changed, 4 insertions(+), 18 deletions(-) diff --git a/docs/contents/usage/fleet.md b/docs/contents/usage/fleet.md index 0c7ae60..4033b89 100644 --- a/docs/contents/usage/fleet.md +++ b/docs/contents/usage/fleet.md @@ -91,31 +91,17 @@ async def _connect(self) -> None: For targeted, event-driven changes (e.g. a robot joining or leaving the fleet while the connector runs), use `add_robot()` and `remove_robot()`: ```python -# A robot appeared — create and connect its session immediately +# A robot appeared: create and connect its session immediately self.add_robot(RobotConfig(robot_id="robot-42", cameras=[])) -# A robot left — disconnect and free its session, clearing its state +# A robot left: disconnect and free its session, clearing its state self.remove_robot("robot-42") ``` -A typical autodiscovery loop diffs the fleet manager's current robot set against `self.robot_ids`: - -```python -@override -async def _execution_loop(self) -> None: - discovered = set(await self._fleet_manager.fetch_robot_list()) - current = set(self.robot_ids) - for robot_id in discovered - current: - self.add_robot(RobotConfig(robot_id=robot_id)) - for robot_id in current - discovered: - self.remove_robot(robot_id) - ... # publish data for self.robot_ids -``` - :::{note} -- These methods create or destroy sessions **immediately**, so call them once the connector is connecting or running (from `_connect()` onward), not before `start()`. +- These methods create or destroy sessions **immediately**. They should be called once the connector is connecting or running (from `_connect()` onwards), not before `start()`. - `add_robot()` raises `ValueError` on a duplicate `robot_id`; `remove_robot()` is a no-op (logs a warning) for an unknown id, so it is safe to call from a discovery loop that may fire repeatedly. -- The fleet may shrink to zero robots at runtime; the execution loop simply idles until robots are discovered again. (The loaded configuration still requires at least one robot at startup.) +- The fleet may shrink to zero robots at runtime; the execution loop simply idles until robots are discovered again. - A robot present in both the old and new fleet is treated as unchanged even if its `RobotConfig` differs (e.g. its cameras changed). To apply a changed config to a running robot, call `remove_robot()` then `add_robot()`. - These methods are thread-safe and may be called from the execution loop, a command handler, or any other thread. ::: From aa8ffc7fac3f28909e41cd55ad1f25e80bfc8c87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Wed, 3 Jun 2026 22:22:02 -0300 Subject: [PATCH 04/11] Allow connectors to start with an empty fleet Make ConnectorRootConfig.fleet optional (default []) and drop the must_contain_at_least_one_robot validator, so a connector can start knowing no robots and discover them at runtime via add_robot()/ update_fleet(). robot_ids_must_be_unique still applies (and is a no-op for an empty list). Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/contents/usage/fleet.md | 2 +- inorbit_connector/models.py | 22 ++++------------------ tests/test_models.py | 16 +++++++++++----- 3 files changed, 16 insertions(+), 24 deletions(-) diff --git a/docs/contents/usage/fleet.md b/docs/contents/usage/fleet.md index 4033b89..3f26db4 100644 --- a/docs/contents/usage/fleet.md +++ b/docs/contents/usage/fleet.md @@ -101,7 +101,7 @@ self.remove_robot("robot-42") :::{note} - These methods create or destroy sessions **immediately**. They should be called once the connector is connecting or running (from `_connect()` onwards), not before `start()`. - `add_robot()` raises `ValueError` on a duplicate `robot_id`; `remove_robot()` is a no-op (logs a warning) for an unknown id, so it is safe to call from a discovery loop that may fire repeatedly. -- The fleet may shrink to zero robots at runtime; the execution loop simply idles until robots are discovered again. +- A connector may start with an empty fleet (omit `fleet` in the config, or pass `[]`) and discover all of its robots at runtime. The fleet may likewise shrink to zero robots while running; the execution loop simply idles until robots are discovered again. - A robot present in both the old and new fleet is treated as unchanged even if its `RobotConfig` differs (e.g. its cameras changed). To apply a changed config to a running robot, call `remove_robot()` then `add_robot()`. - These methods are thread-safe and may be called from the execution loop, a command handler, or any other thread. ::: diff --git a/inorbit_connector/models.py b/inorbit_connector/models.py index 528b49c..448b0fe 100644 --- a/inorbit_connector/models.py +++ b/inorbit_connector/models.py @@ -326,7 +326,9 @@ class ConnectorRootConfig(BaseSettings, Generic[T]): env_vars (dict[str, str], optional): Environment variables to be set in the connector or user scripts. The key is the environment variable name and the value is the value to set. - fleet (list[RobotConfig]): The list of robot configurations. + fleet (list[RobotConfig], optional): The list of robot configurations. + Defaults to an empty list, so a connector may start with no robots and + discover them at runtime via ``add_robot()``/``update_fleet()``. """ # All fields are resolvable from ``INORBIT_`` env vars or from @@ -358,7 +360,7 @@ class ConnectorRootConfig(BaseSettings, Generic[T]): maps: dict[str, MapConfig] = {} env_vars: dict[str, str] = {} metrics: MetricsConfig = MetricsConfig() - fleet: list[RobotConfig] + fleet: list[RobotConfig] = [] def __init__(self, **kwargs): # pydantic-settings consumes _env_file before model validators @@ -462,22 +464,6 @@ def to_singular_config(self, robot_id: str) -> Self: ) return config - @field_validator("fleet") - def must_contain_at_least_one_robot( - cls, fleet: list[RobotConfig] - ) -> list[RobotConfig]: - """Validate that the fleet contains at least one robot. - - Args: - fleet (list[RobotConfig]): The fleet configuration - - Returns: - list[RobotConfig]: The fleet configuration - """ - if len(fleet) < 1: - raise ValueError("Fleet must contain at least one robot") - return fleet - @field_validator("fleet") def robot_ids_must_be_unique(cls, fleet: list[RobotConfig]) -> list[RobotConfig]: """Validate that the robot ids are unique. diff --git a/tests/test_models.py b/tests/test_models.py index a966f85..7e335cd 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -75,13 +75,19 @@ def test_with_valid_input(self, base_model): assert model.fleet[0].robot_id == "robot1" assert model.fleet[1].robot_id == "robot2" - def test_fleet_must_contain_at_least_one_robot(self, base_model): + def test_fleet_may_be_empty(self, base_model): + """An empty fleet is allowed so connectors can discover robots at runtime.""" init_input = base_model.copy() init_input["fleet"] = [] - with pytest.raises( - ValidationError, match="Fleet must contain at least one robot" - ): - ConnectorRootConfig(**init_input, _env_file=None) + model = ConnectorRootConfig(**init_input, _env_file=None) + assert model.fleet == [] + + def test_fleet_defaults_to_empty_when_omitted(self, base_model): + """``fleet`` is optional and defaults to an empty list.""" + init_input = base_model.copy() + del init_input["fleet"] + model = ConnectorRootConfig(**init_input, _env_file=None) + assert model.fleet == [] def test_robot_ids_must_be_unique(self, base_model): init_input = base_model.copy() From 0e0bfc5642d565bdf38f902249ecbf41a4690f5e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Wed, 3 Jun 2026 22:25:51 -0300 Subject: [PATCH 05/11] Update fleet.md --- docs/contents/usage/fleet.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/contents/usage/fleet.md b/docs/contents/usage/fleet.md index 3f26db4..339a4ab 100644 --- a/docs/contents/usage/fleet.md +++ b/docs/contents/usage/fleet.md @@ -101,7 +101,7 @@ self.remove_robot("robot-42") :::{note} - These methods create or destroy sessions **immediately**. They should be called once the connector is connecting or running (from `_connect()` onwards), not before `start()`. - `add_robot()` raises `ValueError` on a duplicate `robot_id`; `remove_robot()` is a no-op (logs a warning) for an unknown id, so it is safe to call from a discovery loop that may fire repeatedly. -- A connector may start with an empty fleet (omit `fleet` in the config, or pass `[]`) and discover all of its robots at runtime. The fleet may likewise shrink to zero robots while running; the execution loop simply idles until robots are discovered again. +- A connector may start with an empty fleet (omit `fleet` in the config, or pass `[]`) and add its robots at runtime. The fleet may likewise shrink to zero robots while running. - A robot present in both the old and new fleet is treated as unchanged even if its `RobotConfig` differs (e.g. its cameras changed). To apply a changed config to a running robot, call `remove_robot()` then `add_robot()`. - These methods are thread-safe and may be called from the execution loop, a command handler, or any other thread. ::: From d13d8c6e8142814f47f28d6c572c1b6280e5222e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Wed, 3 Jun 2026 22:28:06 -0300 Subject: [PATCH 06/11] Update docs --- docs/contents/usage/fleet.md | 6 ++---- inorbit_connector/models.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/docs/contents/usage/fleet.md b/docs/contents/usage/fleet.md index 339a4ab..5359c70 100644 --- a/docs/contents/usage/fleet.md +++ b/docs/contents/usage/fleet.md @@ -77,10 +77,8 @@ This makes it the entry point both for **initial provisioning** (declaring the f @override async def _connect(self) -> None: """Connect to fleet manager and declare the fleet.""" - # Fetch robot list from fleet manager API + # Fetch the robot list from the fleet manager API robots = await self._fleet_manager.get_robots() - - # Reconcile the fleet to the discovered robots self.update_fleet( [RobotConfig(robot_id=robot.id, cameras=robot.cameras) for robot in robots] ) @@ -100,7 +98,7 @@ self.remove_robot("robot-42") :::{note} - These methods create or destroy sessions **immediately**. They should be called once the connector is connecting or running (from `_connect()` onwards), not before `start()`. -- `add_robot()` raises `ValueError` on a duplicate `robot_id`; `remove_robot()` is a no-op (logs a warning) for an unknown id, so it is safe to call from a discovery loop that may fire repeatedly. +- `add_robot()` raises `ValueError` on a duplicate `robot_id`; `remove_robot()` is a no-op (logs a warning) for an unknown id, so it is safe to call from a loop that may fire repeatedly. - A connector may start with an empty fleet (omit `fleet` in the config, or pass `[]`) and add its robots at runtime. The fleet may likewise shrink to zero robots while running. - A robot present in both the old and new fleet is treated as unchanged even if its `RobotConfig` differs (e.g. its cameras changed). To apply a changed config to a running robot, call `remove_robot()` then `add_robot()`. - These methods are thread-safe and may be called from the execution loop, a command handler, or any other thread. diff --git a/inorbit_connector/models.py b/inorbit_connector/models.py index 448b0fe..fbb8b83 100644 --- a/inorbit_connector/models.py +++ b/inorbit_connector/models.py @@ -328,7 +328,7 @@ class ConnectorRootConfig(BaseSettings, Generic[T]): value is the value to set. fleet (list[RobotConfig], optional): The list of robot configurations. Defaults to an empty list, so a connector may start with no robots and - discover them at runtime via ``add_robot()``/``update_fleet()``. + manage them at runtime via ``add_robot()``/``update_fleet()``. """ # All fields are resolvable from ``INORBIT_`` env vars or from From fafdee440c24b57560c8b783586d04da382101ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Wed, 3 Jun 2026 22:28:27 -0300 Subject: [PATCH 07/11] Update test_connector.py --- tests/test_connector.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_connector.py b/tests/test_connector.py index d85e7bf..3020b12 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -842,9 +842,7 @@ def test_add_robot_creates_and_connects_session( assert "R3" in self._sessions(fleet_connector) mock_robot_session_pool.get_session.assert_any_call("R3", robot_name="R3") - def test_add_robot_duplicate_raises( - self, fleet_connector, mock_robot_session_pool - ): + def test_add_robot_duplicate_raises(self, fleet_connector, mock_robot_session_pool): mock_robot_session_pool.get_session.reset_mock() with pytest.raises(ValueError): fleet_connector.add_robot(RobotConfig(robot_id="TestRobot1")) @@ -870,9 +868,7 @@ def test_remove_robot_frees_session_and_clears_state( "TestRobot1" not in fleet_connector._FleetConnector__last_published_frame_ids ) - assert ( - "TestRobot1" not in fleet_connector._FleetConnector__pending_system_stats - ) + assert "TestRobot1" not in fleet_connector._FleetConnector__pending_system_stats def test_remove_robot_unknown_is_noop( self, fleet_connector, mock_robot_session_pool, caplog From a6c2122d5e00aeca44f1ad178ea7775c0d9e19a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Fri, 5 Jun 2026 09:36:20 -0300 Subject: [PATCH 08/11] Refactor: single source of truth for robotIds --- docs/contents/usage/fleet.md | 2 +- examples/simple-fleet-connector/connector.py | 12 +- inorbit_connector/connector.py | 129 +++++----- tests/test_connector.py | 247 +++++++++++++------ 4 files changed, 256 insertions(+), 134 deletions(-) diff --git a/docs/contents/usage/fleet.md b/docs/contents/usage/fleet.md index 5359c70..41cb0e3 100644 --- a/docs/contents/usage/fleet.md +++ b/docs/contents/usage/fleet.md @@ -69,7 +69,7 @@ for robot_id in self.robot_ids: ### Updating the Fleet -`update_fleet()` reconciles the connector's live fleet to exactly the list you passed. It creates and connects a session for each newly added robot, and disconnects and frees the session of each robot that is no longer present. Robots already in the fleet are left untouched. +`update_fleet()` reconciles the connector's live fleet membership to the list passed. It creates and connects a session for each newly added robot, and disconnects and frees the session of each robot that is no longer present. Robots already in the fleet are left untouched and their session is not re-created even if their `RobotConfig` changed. This makes it the entry point both for **initial provisioning** (declaring the fleet from fleet-manager data instead of hardcoded config) and for **runtime autodiscovery**: diff --git a/examples/simple-fleet-connector/connector.py b/examples/simple-fleet-connector/connector.py index 524a504..e7f893e 100644 --- a/examples/simple-fleet-connector/connector.py +++ b/examples/simple-fleet-connector/connector.py @@ -130,14 +130,16 @@ async def _execution_loop(self) -> None: publishing operation. """ + # Snapshot the fleet once: robot_ids may change at runtime where runtime fleet + # updates are implemented, so reading it once ensures consistency. + robot_ids = self.robot_ids + # Fetch data for all robots concurrently - robot_data_tasks = [ - get_fleet_robot_data(robot_id) for robot_id in self.robot_ids - ] + robot_data_tasks = [get_fleet_robot_data(robot_id) for robot_id in robot_ids] robot_data_list = await asyncio.gather(*robot_data_tasks) # Create a mapping of robot_id to data - robot_data_map = dict(zip(self.robot_ids, robot_data_list)) + robot_data_map = dict(zip(robot_ids, robot_data_list)) # Publish data for each robot for robot_id, data in robot_data_map.items(): @@ -167,7 +169,7 @@ async def _execution_loop(self) -> None: self.publish_robot_odometry(robot_id, **odometry) self._logger.info( - f"Fleet data updated and published for {len(self.robot_ids)} robots" + f"Fleet data updated and published for {len(robot_ids)} robots" ) @override diff --git a/inorbit_connector/connector.py b/inorbit_connector/connector.py index 2a4cfc2..226456f 100644 --- a/inorbit_connector/connector.py +++ b/inorbit_connector/connector.py @@ -99,12 +99,6 @@ def __init__(self, config: ConnectorRootConfig, **kwargs) -> None: # Common information self.config = config - # Cache of robot IDs in config.fleet. Accessed through the robot_ids property. - # Seeded directly from the config here (no sessions are created during - # construction); thereafter kept in sync by update_fleet()/add_robot()/ - # remove_robot(). The initial robot sessions are created later by __connect(), - # which calls update_fleet(self.config.fleet) once the connector thread is up. - self.__robot_ids: list[str] = [robot.robot_id for robot in config.fleet] # Per robot state self.__last_published_frame_ids: dict[str, str] = {} @@ -120,8 +114,8 @@ def __init__(self, config: ConnectorRootConfig, **kwargs) -> None: # robot_id) for thread-safe access) in tight loops. It should not be accessed # directly by subclasses to maintain thread-safety self.__robot_sessions: dict[str, RobotSession] = {} - # Lock guarding fleet mutation: __robot_ids, __robot_sessions membership, - # config.fleet, and the per-robot state dicts. + # Lock guarding fleet mutation: config.fleet, __robot_sessions membership, and + # the per-robot state dicts (last-frame-id, pending system stats). self.__fleet_lock = threading.RLock() # Threading for the main run methods @@ -213,10 +207,9 @@ def __init__(self, config: ConnectorRootConfig, **kwargs) -> None: @property def robot_ids(self) -> list[str]: """Get the list of robot IDs in the fleet.""" - # Return a copy of the cached list so callers can't mutate internal state and - # so a captured reference stays consistent if the fleet changes concurrently. + # Return a copy of the fleet robotIds derived from config.fleet under the lock. with self.__fleet_lock: - return self.__robot_ids.copy() + return [robot.robot_id for robot in self.config.fleet] @property def _connector_type(self) -> str: @@ -240,11 +233,12 @@ def update_fleet(self, fleet: list[RobotConfig]) -> None: - disconnects and frees the session of each robot no longer present, clearing its per-robot state. - It is idempotent: robots already in the fleet are left untouched. - ``__connect()`` calls it with ``self.config.fleet`` to create the - initial sessions, and subclasses may call it at runtime (from ``_connect()`` - onward) to implement robot autodiscovery. e.g. fetching the robot list from a - fleet manager API. + It is idempotent: robots already in the fleet are left untouched. It reconciles + membership only. A robot present in both the old and new fleet keeps its + existing session even if its ``RobotConfig`` changed. The connector's startup + reconciles ``self.config.fleet`` to create the initial sessions, and subclasses + may call this at runtime (from ``_connect()`` onward) to implement fleet runtime + updates, e.g. fetching the robot list from a fleet manager API. Note: Sessions are created/destroyed immediately, so this must be called once the @@ -280,15 +274,27 @@ def __apply_fleet(self, fleet: list[RobotConfig]) -> None: # respect to concurrent add_robot()/remove_robot()/update_fleet() calls. The # lock is re-entrant, so the wrappers can hold it across their read+delegate. with self.__fleet_lock: - new_id_set = set(new_ids) - to_remove = [rid for rid in self.__robot_sessions if rid not in new_id_set] + configs = {robot.robot_id: robot for robot in fleet} + to_remove = [rid for rid in self.__robot_sessions if rid not in configs] to_add = [rid for rid in new_ids if rid not in self.__robot_sessions] - # Apply the new membership before touching sessions so the running loop - # sees the target fleet and stops iterating removed robots. + # Create the new sessions before mutating any membership or state, so a + # failure to connect rolls back cleanly and never leaves the fleet listing a + # robot without a session. + created: dict[str, RobotSession] = {} + try: + for robot_id in to_add: + created[robot_id] = self.__initialize_session(configs[robot_id]) + except Exception: + # get_session() registers the session in the pool before connecting, so + # free every attempted add (a no-op when the pool has none) to undo any + # half-built session, then re-raise with membership untouched. + for robot_id in to_add: + self.__session_pool.free_robot_session(robot_id) + raise + + # Commit the membership and drop removed robots. self.config.fleet = list(fleet) - self.__robot_ids = list(new_ids) - for robot_id in to_remove: self.__last_published_frame_ids.pop(robot_id, None) self.__pending_system_stats.pop(robot_id, None) @@ -296,8 +302,7 @@ def __apply_fleet(self, fleet: list[RobotConfig]) -> None: # free_robot_session disconnects and removes the session from the # pool; no-op if the pool has no session for this robot. self.__session_pool.free_robot_session(robot_id) - for robot_id in to_add: - self.__robot_sessions[robot_id] = self.__initialize_session(robot_id) + self.__robot_sessions.update(created) def add_robot(self, robot_config: RobotConfig) -> None: """Add a single robot to the fleet at runtime. @@ -317,7 +322,7 @@ def add_robot(self, robot_config: RobotConfig) -> None: ValueError: If a robot with the same ``robot_id`` is already in the fleet. """ with self.__fleet_lock: - if robot_config.robot_id in self.__robot_ids: + if any(r.robot_id == robot_config.robot_id for r in self.config.fleet): raise ValueError( f"Robot '{robot_config.robot_id}' is already in the fleet" ) @@ -335,7 +340,7 @@ def remove_robot(self, robot_id: str) -> None: robot_id (str): The ID of the robot to remove. """ with self.__fleet_lock: - if robot_id not in self.__robot_ids: + if not any(r.robot_id == robot_id for r in self.config.fleet): self._logger.warning( f"remove_robot: robot '{robot_id}' is not in the fleet" ) @@ -448,9 +453,10 @@ def __register_user_scripts_for_session( # bash scripts session.register_commands_path(path, exec_name_regex=r".*\.sh") - def __initialize_session(self, robot_id: str) -> RobotSession: - """Initialize a robot session.""" + def __initialize_session(self, robot_config: RobotConfig) -> RobotSession: + """Initialize a robot session for the given robot configuration.""" + robot_id = robot_config.robot_id # The InOrbit hostname of the robot is set to the robot_id by default # There is no support for setting the display name of a robot through the # edge-sdk yet @@ -477,18 +483,16 @@ def __initialize_session(self, robot_id: str) -> RobotSession: ) # Register cameras declared in this robot's configuration. - robot_config = self._get_robot_config(robot_id) - if robot_config is not None: - for idx, camera_config in enumerate(robot_config.cameras): - self._logger.info( - f"Registering camera {idx} for robot {robot_id}: " - f"{str(camera_config.video_url)}" - ) - # If values are None, remove the key from the dictionary to use - # edge-sdk defaults - dump = camera_config.model_dump() - clean = {k: v for k, v in dump.items() if v is not None} - session.register_camera(str(idx), OpenCVCamera(**clean)) + for idx, camera_config in enumerate(robot_config.cameras): + self._logger.info( + f"Registering camera {idx} for robot {robot_id}: " + f"{str(camera_config.video_url)}" + ) + # If values are None, remove the key from the dictionary to use + # edge-sdk defaults + dump = camera_config.model_dump() + clean = {k: v for k, v in dump.items() if v is not None} + session.register_camera(str(idx), OpenCVCamera(**clean)) return session @@ -515,13 +519,16 @@ async def __connect(self) -> None: async def __disconnect(self) -> None: """Disconnect external services and disconnect from InOrbit.""" - # Disconnect from InOrbit while holding the lock so no robot can be added (via - # add_robot/update_fleet from another thread) while we tear the fleet down, and - # clear the session dict so a subsequent start() recreates the sessions. + # Tear the fleet down while holding the lock so no robot can be added (via + # add_robot/update_fleet from another thread) mid-teardown. Free each session + # from the pool (not just disconnect it) so a subsequent start() rebuilds and + # reconnects fresh sessions instead of reusing stale, disconnected ones. with self.__fleet_lock: - for session in self.__robot_sessions.values(): - session.disconnect() + for robot_id in list(self.__robot_sessions): + self.__session_pool.free_robot_session(robot_id) self.__robot_sessions.clear() + self.__last_published_frame_ids.clear() + self.__pending_system_stats.clear() # Clean up temporary map files if self.__temp_map_dir is not None: @@ -578,7 +585,8 @@ async def __run_loop(self) -> None: _metrics.execution_loop_errors.add(1) self._logger.error(f"Error in execution loop: {e}") self._logger.error(f"Traceback: {traceback.format_exc()}") - self.__pending_system_stats.clear() + with self.__fleet_lock: + self.__pending_system_stats.clear() # Continue execution after a brief pause to avoid tight error loops await asyncio.sleep(1.0) @@ -588,8 +596,9 @@ def _get_robot_session(self, robot_id: str) -> RobotSession: Usually the connector API is enough to abstract from the edge-sdk, but in some cases accessing the robot session directly may be necessary. - This method provides thread-safe access to robot sessions through the session - pool. + This method provides thread-safe access to the connector's active robot + sessions. It reads the connector's own session map rather than the pool, so it + never creates a session for a robot that is not in the fleet. Args: robot_id (str): The robot ID to get the session for @@ -598,9 +607,12 @@ def _get_robot_session(self, robot_id: str) -> RobotSession: RobotSession: The robot session for the specified robot Raises: - KeyError: If the robot_id is not found in the pool + KeyError: If the robot_id is not in the fleet """ - return self.__session_pool.get_session(robot_id) + session = self.__robot_sessions.get(robot_id) + if session is None: + raise KeyError(f"No active session for robot '{robot_id}'") + return session def _is_session_connected(self, robot_id: str) -> bool: """Return True if the MQTT session for ``robot_id`` is currently connected. @@ -694,12 +706,15 @@ def publish_robot_pose( **kwargs: Additional arguments for pose publishing """ session = self._get_robot_session(robot_id) - last_published_frame_id = self.__last_published_frame_ids.get(robot_id, None) - if frame_id != last_published_frame_id: + with self.__fleet_lock: + changed = frame_id != self.__last_published_frame_ids.get(robot_id) + if changed: + self.__last_published_frame_ids[robot_id] = frame_id + if changed: self._logger.info( f"Updating map {frame_id} with new pose for robot {robot_id}." ) - self.__last_published_frame_ids[robot_id] = frame_id + # map/pose publish is I/O and should not hold the fleet lock self.publish_robot_map(robot_id, frame_id, is_update=True) session.publish_pose(x, y, yaw, frame_id, **kwargs) @@ -925,11 +940,7 @@ def __publish_pending_system_stats(self) -> None: ) with self.__fleet_lock: - sessions = { - robot_id: self.__robot_sessions[robot_id] - for robot_id in self.__robot_ids - if robot_id in self.__robot_sessions - } + sessions = dict(self.__robot_sessions) pending = self.__pending_system_stats self.__pending_system_stats = {} @@ -1040,7 +1051,7 @@ class Connector(FleetConnector, ABC): @override def update_fleet(self, fleet: list[RobotConfig]) -> None: raise NotImplementedError( - "udpate_fleet() is not supported on a single-robot Connector; " + "update_fleet() is not supported on a single-robot Connector; " "use FleetConnector for multi-robot fleets" ) diff --git a/tests/test_connector.py b/tests/test_connector.py index 3020b12..77ee1aa 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -76,6 +76,21 @@ def get_session_side_effect(robot_id, robot_name=None): yield mock_pool +def _seed_sessions(connector): + """Mimic the post-``__connect()`` state for tests that don't call ``start()``. + + Production creates the per-robot sessions in ``__connect()``; tests that build a + connector directly must seed ``__robot_sessions`` so the ``publish_*`` / + ``_get_robot_session`` paths (which read the connector's own session map) work. + """ + pool = connector._FleetConnector__session_pool + connector._FleetConnector__robot_sessions = { + rc.robot_id: pool.get_session(rc.robot_id, robot_name=rc.robot_id) + for rc in connector.config.fleet + } + return connector + + # ============================================================================== # FleetConnector Tests # ============================================================================== @@ -147,17 +162,22 @@ def base_model(self): @pytest.fixture(autouse=True) def make_fleet_connector_not_abstract(self): - FleetConnector.__abstractmethods__ = set() + saved = FleetConnector.__abstractmethods__ + FleetConnector.__abstractmethods__ = frozenset() + yield + FleetConnector.__abstractmethods__ = saved @pytest.fixture def base_fleet_connector(self, base_model): - return FleetConnector( - ConnectorRootConfig( - **base_model, - fleet=[ - RobotConfig(robot_id="TestRobot1"), - RobotConfig(robot_id="TestRobot2"), - ], + return _seed_sessions( + FleetConnector( + ConnectorRootConfig( + **base_model, + fleet=[ + RobotConfig(robot_id="TestRobot1"), + RobotConfig(robot_id="TestRobot2"), + ], + ) ) ) @@ -237,9 +257,7 @@ def test_publish_robot_pose(self, base_fleet_connector, mock_robot_session_pool) robot_id = "TestRobot1" base_fleet_connector.publish_robot_pose(robot_id, 1.0, 2.0, 3.14, "map1") - # Verify get_session was called and get the actual session returned - mock_robot_session_pool.get_session.assert_called_with(robot_id) - # Get the session from the side_effect cache + # The robot's active session received the pose. session = base_fleet_connector._get_robot_session(robot_id) session.publish_pose.assert_called() @@ -255,10 +273,12 @@ def test_publish_robot_pose_updates_map(self, base_model, mock_robot_session_poo "resolution": 0.1, } } - connector = FleetConnector( - ConnectorRootConfig( - **base_model, - fleet=[RobotConfig(robot_id="TestRobot1")], + connector = _seed_sessions( + FleetConnector( + ConnectorRootConfig( + **base_model, + fleet=[RobotConfig(robot_id="TestRobot1")], + ) ) ) @@ -412,14 +432,19 @@ def base_model(self): @pytest.fixture(autouse=True) def make_fleet_connector_not_abstract(self): - FleetConnector.__abstractmethods__ = set() + saved = FleetConnector.__abstractmethods__ + FleetConnector.__abstractmethods__ = frozenset() + yield + FleetConnector.__abstractmethods__ = saved @pytest.fixture def fleet_connector(self, base_model): - return FleetConnector( - ConnectorRootConfig( - **base_model, - fleet=[RobotConfig(robot_id="TestRobot1")], + return _seed_sessions( + FleetConnector( + ConnectorRootConfig( + **base_model, + fleet=[RobotConfig(robot_id="TestRobot1")], + ) ) ) @@ -591,7 +616,10 @@ def base_model(self): @pytest.fixture(autouse=True) def make_fleet_connector_not_abstract(self): - FleetConnector.__abstractmethods__ = set() + saved = FleetConnector.__abstractmethods__ + FleetConnector.__abstractmethods__ = frozenset() + yield + FleetConnector.__abstractmethods__ = saved @pytest.fixture def fleet_connector(self, base_model, mock_robot_session_pool): @@ -807,7 +835,10 @@ def base_model(self): @pytest.fixture(autouse=True) def make_fleet_connector_not_abstract(self): - FleetConnector.__abstractmethods__ = set() + saved = FleetConnector.__abstractmethods__ + FleetConnector.__abstractmethods__ = frozenset() + yield + FleetConnector.__abstractmethods__ = saved @pytest.fixture def fleet_connector(self, base_model, mock_robot_session_pool): @@ -879,6 +910,18 @@ def test_remove_robot_unknown_is_noop( assert fleet_connector.robot_ids == ["TestRobot1", "TestRobot2"] assert "not in the fleet" in caplog.text + def test_publish_to_removed_robot_does_not_resurrect_session( + self, fleet_connector, mock_robot_session_pool + ): + """Accessing/publishing to a removed robot must raise, not silently recreate a + zombie session via the pool.""" + fleet_connector.remove_robot("TestRobot1") + + with pytest.raises(KeyError): + fleet_connector._get_robot_session("TestRobot1") + with pytest.raises(KeyError): + fleet_connector.publish_robot_odometry("TestRobot1", linear_speed=1.0) + def test_update_fleet_reconciles_diff( self, fleet_connector, mock_robot_session_pool ): @@ -905,6 +948,44 @@ def test_update_fleet_duplicate_ids_raises(self, fleet_connector): [RobotConfig(robot_id="R3"), RobotConfig(robot_id="R3")] ) + def test_update_fleet_rolls_back_on_session_failure( + self, fleet_connector, mock_robot_session_pool + ): + """If a session fails to connect mid-reconcile, the whole change rolls back: + membership is untouched and every attempted add is freed from the pool.""" + before_ids = fleet_connector.robot_ids + before_fleet = list(fleet_connector.config.fleet) + + # Fail when creating the 2nd newly-added robot's session. + def failing_get_session(robot_id, robot_name=None): + if robot_id == "R4": + raise RuntimeError("connect failed") + session = MagicMock(spec=RobotSession) + session.robot_id = robot_id + return session + + mock_robot_session_pool.get_session.side_effect = failing_get_session + + with pytest.raises(RuntimeError): + fleet_connector.update_fleet( + [ + RobotConfig(robot_id="TestRobot1"), + RobotConfig(robot_id="TestRobot2"), + RobotConfig(robot_id="R3"), + RobotConfig(robot_id="R4"), + ] + ) + + # Nothing committed: ids, config.fleet and the session map are as before. + assert fleet_connector.robot_ids == before_ids + assert fleet_connector.config.fleet == before_fleet + assert set(self._sessions(fleet_connector)) == set(before_ids) + # Both attempted adds were freed from the pool (incl. the half-built one). + freed = { + c.args[0] for c in mock_robot_session_pool.free_robot_session.call_args_list + } + assert {"R3", "R4"} <= freed + def test_connect_materializes_initial_sessions( self, base_model, mock_robot_session_pool ): @@ -934,7 +1015,8 @@ def test_connect_materializes_initial_sessions( def test_disconnect_clears_sessions_for_restart( self, base_model, mock_robot_session_pool ): - """stop() disconnects + clears sessions so a later start() reconnects.""" + """stop() frees the pool sessions + clears the map so a later start() + rebuilds and reconnects instead of reusing stale, disconnected sessions.""" connector = FleetConnector( ConnectorRootConfig( **base_model, @@ -945,13 +1027,14 @@ def test_disconnect_clears_sessions_for_restart( sleep(0.5) connector.stop() - # Session was created and then disconnected, and the dict cleared. - session = mock_robot_session_pool.get_session("TestRobot1") - session.disconnect.assert_called() + # The session was freed from the pool (which disconnects it) and the map + # cleared. Freeing the pool — not just disconnecting — is what lets a restart + # rebuild a fresh session (get_session skips connect() on a pool cache hit). + mock_robot_session_pool.free_robot_session.assert_any_call("TestRobot1") assert connector._FleetConnector__robot_sessions == {} # Restart must recreate the session (to_add keys off __robot_sessions). - session.reset_mock() + mock_robot_session_pool.free_robot_session.reset_mock() connector.start() try: sleep(0.5) @@ -1013,24 +1096,35 @@ def test_get_robot_config(self, fleet_connector): assert fleet_connector._get_robot_config("nope") is None def test_thread_safety_smoke(self, fleet_connector): - """Concurrent add/remove on disjoint ids leaves consistent state.""" - - def worker(start): - for i in range(start, start + 20): - rid = f"T{i}" - fleet_connector.add_robot(RobotConfig(robot_id=rid)) - fleet_connector.remove_robot(rid) - - threads = [ - threading.Thread(target=worker, args=(base,)) - for base in (100, 200, 300, 400) - ] + """Concurrent add/remove on a SHARED id set must keep the fleet's internal + structures consistent. The shared ids force real contention on __fleet_lock — + without it the read-modify-write in add/remove would corrupt config.fleet or + desync it from the session map. Duplicate adds racing each other are expected + and raise ValueError, which the workers swallow.""" + shared_ids = [f"T{i}" for i in range(8)] + + def worker(): + for _ in range(50): + for rid in shared_ids: + try: + fleet_connector.add_robot(RobotConfig(robot_id=rid)) + except ValueError: + # Another thread won the race to add this id — expected. + pass + fleet_connector.remove_robot(rid) # no-op if already gone + + threads = [threading.Thread(target=worker) for _ in range(4)] for t in threads: t.start() for t in threads: t.join() - # All transient robots removed; the original fleet remains; no duplicates. + # Deterministic cleanup so the final assertion doesn't hinge on interleaving. + for rid in shared_ids: + fleet_connector.remove_robot(rid) + + # The original fleet remains, with no duplicates, and robot_ids agrees with the + # live session map (the invariant the lock protects). ids = fleet_connector.robot_ids assert sorted(ids) == ["TestRobot1", "TestRobot2"] assert len(ids) == len(set(ids)) @@ -1105,15 +1199,20 @@ def base_model(self): @pytest.fixture(autouse=True) def make_connector_not_abstract(self): - Connector.__abstractmethods__ = set() + saved = Connector.__abstractmethods__ + Connector.__abstractmethods__ = frozenset() + yield + Connector.__abstractmethods__ = saved @pytest.fixture def base_connector(self, base_model): - return Connector( - "TestRobot", - ConnectorRootConfig( - **base_model, fleet=[RobotConfig(robot_id="TestRobot")] - ), + return _seed_sessions( + Connector( + "TestRobot", + ConnectorRootConfig( + **base_model, fleet=[RobotConfig(robot_id="TestRobot")] + ), + ) ) def test_init(self, base_model): @@ -1141,7 +1240,7 @@ def test_init_with_robot_key(self, base_model, mock_robot_session_pool): ) robot_id = "TestRobot" - connector = Connector(robot_id, config) + connector = _seed_sessions(Connector(robot_id, config)) session = connector._get_session() assert session.robot_id == "TestRobot" @@ -1179,7 +1278,6 @@ def test_start_still_creates_session_despite_blocked_update_fleet( ) finally: connector.stop() - mock_robot_session_pool.get_session.assert_called() def test_publish_map(self, base_model, mock_robot_session_pool): """Test publish_map delegates to FleetConnector.publish_robot_map.""" @@ -1193,11 +1291,13 @@ def test_publish_map(self, base_model, mock_robot_session_pool): "resolution": 0.1, } } - connector = Connector( - "TestRobot", - ConnectorRootConfig( - **base_model, fleet=[RobotConfig(robot_id="TestRobot")] - ), + connector = _seed_sessions( + Connector( + "TestRobot", + ConnectorRootConfig( + **base_model, fleet=[RobotConfig(robot_id="TestRobot")] + ), + ) ) connector.publish_map("frameA") @@ -1217,11 +1317,13 @@ def test_publish_pose(self, base_model, mock_robot_session_pool): "resolution": 0.1, } } - connector = Connector( - "TestRobot", - ConnectorRootConfig( - **base_model, fleet=[RobotConfig(robot_id="TestRobot")] - ), + connector = _seed_sessions( + Connector( + "TestRobot", + ConnectorRootConfig( + **base_model, fleet=[RobotConfig(robot_id="TestRobot")] + ), + ) ) connector.publish_pose(1.0, 2.0, 3.14, "frameA") @@ -1248,11 +1350,13 @@ def test_publish_pose_updates_maps(self, base_model, mock_robot_session_pool): "resolution": 0.1, }, } - connector = Connector( - "TestRobot", - ConnectorRootConfig( - **base_model, fleet=[RobotConfig(robot_id="TestRobot")] - ), + connector = _seed_sessions( + Connector( + "TestRobot", + ConnectorRootConfig( + **base_model, fleet=[RobotConfig(robot_id="TestRobot")] + ), + ) ) session = connector._get_session() @@ -1434,15 +1538,20 @@ def base_model(self): @pytest.fixture(autouse=True) def make_connector_not_abstract(self): - Connector.__abstractmethods__ = set() + saved = Connector.__abstractmethods__ + Connector.__abstractmethods__ = frozenset() + yield + Connector.__abstractmethods__ = saved @pytest.fixture def base_connector(self, base_model): - return Connector( - "TestRobot", - ConnectorRootConfig( - **base_model, fleet=[RobotConfig(robot_id="TestRobot")] - ), + return _seed_sessions( + Connector( + "TestRobot", + ConnectorRootConfig( + **base_model, fleet=[RobotConfig(robot_id="TestRobot")] + ), + ) ) @pytest.mark.asyncio From 2d56960226be95683ab3118686ca92a6e4e779f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Fri, 5 Jun 2026 10:19:40 -0300 Subject: [PATCH 09/11] Remove extra method --- inorbit_connector/connector.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/inorbit_connector/connector.py b/inorbit_connector/connector.py index 226456f..8b62614 100644 --- a/inorbit_connector/connector.py +++ b/inorbit_connector/connector.py @@ -349,22 +349,6 @@ def remove_robot(self, robot_id: str) -> None: [robot for robot in self.config.fleet if robot.robot_id != robot_id] ) - def _get_robot_config(self, robot_id: str) -> RobotConfig | None: - """Return the ``RobotConfig`` for ``robot_id`` from the fleet, or None. - - Args: - robot_id (str): The robot ID to look up. - - Returns: - RobotConfig | None: The matching configuration, or None if the robot is not - in the fleet. - """ - with self.__fleet_lock: - for robot_config in self.config.fleet: - if robot_config.robot_id == robot_id: - return robot_config - return None - def _handle_command_exception( self, exception: Exception, From dd01c0194e84b6c6d2ff864ae8ecbf3df3d70b67 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Fri, 5 Jun 2026 11:33:34 -0300 Subject: [PATCH 10/11] Fix a possible publish block --- inorbit_connector/connector.py | 93 ++++++++++---- inorbit_connector/models.py | 6 +- tests/test_connector.py | 223 +++++++++++++++++++++++++-------- 3 files changed, 245 insertions(+), 77 deletions(-) diff --git a/inorbit_connector/connector.py b/inorbit_connector/connector.py index 8b62614..c0213bd 100644 --- a/inorbit_connector/connector.py +++ b/inorbit_connector/connector.py @@ -256,6 +256,12 @@ def update_fleet(self, fleet: list[RobotConfig]) -> None: Raises: ValueError: If ``fleet`` contains duplicate robot IDs. """ + + # TODO(b-Tomas): Session creation/teardown (MQTT connect/disconnect) runs while + # holding the fleet lock, so a slow broker connection briefly blocks publishing + # from other threads. Fleet mutations are expected to be infrequent, but + # consider fixing this. + self.__apply_fleet(fleet) def __apply_fleet(self, fleet: list[RobotConfig]) -> None: @@ -293,8 +299,16 @@ def __apply_fleet(self, fleet: list[RobotConfig]) -> None: self.__session_pool.free_robot_session(robot_id) raise - # Commit the membership and drop removed robots. - self.config.fleet = list(fleet) + # Commit the membership and drop removed robots. Kept robots retain their + # existing RobotConfig (their session is not re-created), so config.fleet + # stays consistent with the live sessions; only newly added robots take + # the incoming config. To apply a changed config to a running robot, call + # remove_robot() then add_robot(). + old_by_id = {rc.robot_id: rc for rc in self.config.fleet} + self.config.fleet = [ + rc if rc.robot_id in created else old_by_id.get(rc.robot_id, rc) + for rc in fleet + ] for robot_id in to_remove: self.__last_published_frame_ids.pop(robot_id, None) self.__pending_system_stats.pop(robot_id, None) @@ -326,7 +340,7 @@ def add_robot(self, robot_config: RobotConfig) -> None: raise ValueError( f"Robot '{robot_config.robot_id}' is already in the fleet" ) - self.update_fleet([*self.config.fleet, robot_config]) + self.__apply_fleet([*self.config.fleet, robot_config]) def remove_robot(self, robot_id: str) -> None: """Remove a single robot from the fleet at runtime. @@ -345,7 +359,7 @@ def remove_robot(self, robot_id: str) -> None: f"remove_robot: robot '{robot_id}' is not in the fleet" ) return - self.update_fleet( + self.__apply_fleet( [robot for robot in self.config.fleet if robot.robot_id != robot_id] ) @@ -499,6 +513,11 @@ async def __connect(self) -> None: # directly so it also works on the single-robot Connector subclass, which blocks # the public update_fleet(). self.__apply_fleet(self.config.fleet) + robot_ids = self.robot_ids + self._logger.info( + f"Initialized {len(robot_ids)} robot session(s)" + + (f" for robots {', '.join(robot_ids)}" if robot_ids else "") + ) async def __disconnect(self) -> None: """Disconnect external services and disconnect from InOrbit.""" @@ -574,29 +593,25 @@ async def __run_loop(self) -> None: # Continue execution after a brief pause to avoid tight error loops await asyncio.sleep(1.0) - def _get_robot_session(self, robot_id: str) -> RobotSession: - """Get a robot session for a specific robot ID. + def _get_robot_session(self, robot_id: str) -> RobotSession | None: + """Get the active robot session for a specific robot ID, or None. Usually the connector API is enough to abstract from the edge-sdk, but in some cases accessing the robot session directly may be necessary. This method provides thread-safe access to the connector's active robot sessions. It reads the connector's own session map rather than the pool, so it - never creates a session for a robot that is not in the fleet. + never creates a session for a robot that is not in the fleet, and returns None + when the robot has no active session so the ``publish_*`` paths can skip it. Args: robot_id (str): The robot ID to get the session for Returns: - RobotSession: The robot session for the specified robot - - Raises: - KeyError: If the robot_id is not in the fleet + RobotSession | None: The robot's active session, or None if it has none. """ - session = self.__robot_sessions.get(robot_id) - if session is None: - raise KeyError(f"No active session for robot '{robot_id}'") - return session + with self.__fleet_lock: + return self.__robot_sessions.get(robot_id) def _is_session_connected(self, robot_id: str) -> bool: """Return True if the MQTT session for ``robot_id`` is currently connected. @@ -609,7 +624,7 @@ def _is_session_connected(self, robot_id: str) -> bool: has not been initialized yet, which is the normal state before ``_connect()`` runs. """ - session = self.__robot_sessions.get(robot_id) + session = self._get_robot_session(robot_id) if session is None: return False try: @@ -628,7 +643,8 @@ def start(self) -> None: It: - calls self._connect() to connect to any external services. - - sets up camera feeds defined in the configuration. + - creates and connects the robot sessions for the configured fleet + (registering each robot's cameras, command handler and status callback). - runs the execution loop in a new thread. - calls self._disconnect() to disconnect from any external services once the connector is stopped. @@ -690,6 +706,11 @@ def publish_robot_pose( **kwargs: Additional arguments for pose publishing """ session = self._get_robot_session(robot_id) + if session is None: + self._logger.debug( + f"Skipping pose publish for '{robot_id}': no active session" + ) + return with self.__fleet_lock: changed = frame_id != self.__last_published_frame_ids.get(robot_id) if changed: @@ -699,11 +720,15 @@ def publish_robot_pose( f"Updating map {frame_id} with new pose for robot {robot_id}." ) # map/pose publish is I/O and should not hold the fleet lock - self.publish_robot_map(robot_id, frame_id, is_update=True) + self.publish_robot_map(robot_id, frame_id, is_update=True, session=session) session.publish_pose(x, y, yaw, frame_id, **kwargs) def publish_robot_map( - self, robot_id: str, frame_id: str, is_update: bool = False + self, + robot_id: str, + frame_id: str, + is_update: bool = False, + session: RobotSession | None = None, ) -> None: """Publish the map metadata for a specific robot to InOrbit. @@ -715,8 +740,17 @@ def publish_robot_map( robot_id (str): The robot ID to publish map for frame_id (str): The frame ID of the map is_update (bool): Whether this is an update to an existing map + session (RobotSession | None): The robot's session, if already resolved by + the caller. When omitted it is resolved here; a robot removed + concurrently is skipped rather than raising. """ - session = self._get_robot_session(robot_id) + if session is None: + session = self._get_robot_session(robot_id) + if session is None: + self._logger.debug( + f"Skipping map publish for '{robot_id}': no active session" + ) + return if map_config := self.config.maps.get(frame_id): session.publish_map( file=map_config.file, @@ -849,6 +883,11 @@ def publish_robot_odometry(self, robot_id: str, **kwargs) -> None: **kwargs: Odometry data """ session = self._get_robot_session(robot_id) + if session is None: + self._logger.debug( + f"Skipping odometry publish for '{robot_id}': no active session" + ) + return session.publish_odometry(**kwargs) def publish_robot_key_values(self, robot_id: str, **kwargs) -> None: @@ -863,6 +902,11 @@ def publish_robot_key_values(self, robot_id: str, **kwargs) -> None: **kwargs: Key-value data """ session = self._get_robot_session(robot_id) + if session is None: + self._logger.debug( + f"Skipping key-values publish for '{robot_id}': no active session" + ) + return session.publish_key_values({"connector_type": self._connector_type, **kwargs}) def publish_robot_system_stats(self, robot_id: str, **kwargs) -> None: @@ -1084,8 +1128,15 @@ def _get_session(self) -> RobotSession: Usually the connector API is enough to abstract from the edge-sdk, but in some cases accessing the robot session directly may be necessary. + + Raises: + KeyError: If the robot has no active session (e.g. accessed before the + connector has connected, or after it stopped). """ - return super()._get_robot_session(self.robot_id) + session = super()._get_robot_session(self.robot_id) + if session is None: + raise KeyError(f"No active session for robot '{self.robot_id}'") + return session def _is_robot_online(self) -> bool: """Check if the robot is online. diff --git a/inorbit_connector/models.py b/inorbit_connector/models.py index fbb8b83..68d8faf 100644 --- a/inorbit_connector/models.py +++ b/inorbit_connector/models.py @@ -327,8 +327,10 @@ class ConnectorRootConfig(BaseSettings, Generic[T]): connector or user scripts. The key is the environment variable name and the value is the value to set. fleet (list[RobotConfig], optional): The list of robot configurations. - Defaults to an empty list, so a connector may start with no robots and - manage them at runtime via ``add_robot()``/``update_fleet()``. + Defaults to an empty list. A ``FleetConnector`` may start with no robots + and manage them at runtime via ``add_robot()``/``update_fleet()``. A + single-robot ``Connector`` still requires exactly one matching robot + (``to_singular_config`` raises otherwise). """ # All fields are resolvable from ``INORBIT_`` env vars or from diff --git a/tests/test_connector.py b/tests/test_connector.py index 77ee1aa..0269462 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -6,8 +6,10 @@ # SPDX-License-Identifier: MIT # Standard +import logging import os import threading +from contextlib import contextmanager from time import sleep from unittest.mock import AsyncMock, MagicMock, patch @@ -91,6 +93,20 @@ def _seed_sessions(connector): return connector +@contextmanager +def _concrete(cls): + """Temporarily clear ``cls.__abstractmethods__`` so it can be instantiated. + + Restores the original set on exit so the change does not leak between tests. + """ + saved = cls.__abstractmethods__ + cls.__abstractmethods__ = frozenset() + try: + yield + finally: + cls.__abstractmethods__ = saved + + # ============================================================================== # FleetConnector Tests # ============================================================================== @@ -162,10 +178,8 @@ def base_model(self): @pytest.fixture(autouse=True) def make_fleet_connector_not_abstract(self): - saved = FleetConnector.__abstractmethods__ - FleetConnector.__abstractmethods__ = frozenset() - yield - FleetConnector.__abstractmethods__ = saved + with _concrete(FleetConnector): + yield @pytest.fixture def base_fleet_connector(self, base_model): @@ -432,10 +446,8 @@ def base_model(self): @pytest.fixture(autouse=True) def make_fleet_connector_not_abstract(self): - saved = FleetConnector.__abstractmethods__ - FleetConnector.__abstractmethods__ = frozenset() - yield - FleetConnector.__abstractmethods__ = saved + with _concrete(FleetConnector): + yield @pytest.fixture def fleet_connector(self, base_model): @@ -616,10 +628,8 @@ def base_model(self): @pytest.fixture(autouse=True) def make_fleet_connector_not_abstract(self): - saved = FleetConnector.__abstractmethods__ - FleetConnector.__abstractmethods__ = frozenset() - yield - FleetConnector.__abstractmethods__ = saved + with _concrete(FleetConnector): + yield @pytest.fixture def fleet_connector(self, base_model, mock_robot_session_pool): @@ -632,13 +642,8 @@ def fleet_connector(self, base_model, mock_robot_session_pool): ], ) ) - # Simulate the post-connect state: sessions are created by update_fleet() - # during __connect(). __publish_pending_system_stats reads from this dict. - connector._FleetConnector__robot_sessions = { - robot_id: mock_robot_session_pool.get_session(robot_id) - for robot_id in connector.robot_ids - } - return connector + # Seed the post-connect state (sessions are created in __connect()). + return _seed_sessions(connector) def test_publish_pending_system_stats_publishes_stored_stats( self, fleet_connector, mock_robot_session_pool @@ -747,10 +752,8 @@ def test_publish_connector_system_stats_uses_psutil( ), publish_connector_system_stats=True, ) - # Simulate the post-connect state (sessions created by update_fleet). - connector._FleetConnector__robot_sessions = { - "TestRobot1": mock_robot_session_pool.get_session("TestRobot1") - } + # Seed the post-connect state (sessions created in __connect()). + _seed_sessions(connector) connector._FleetConnector__publish_pending_system_stats() @@ -773,10 +776,8 @@ def test_publish_connector_system_stats_fallback_without_psutil( ), publish_connector_system_stats=True, ) - # Simulate the post-connect state (sessions created by update_fleet). - connector._FleetConnector__robot_sessions = { - "TestRobot1": mock_robot_session_pool.get_session("TestRobot1") - } + # Seed the post-connect state (sessions created in __connect()). + _seed_sessions(connector) # Should have fallen back to False assert connector._FleetConnector__publish_connector_system_stats is False @@ -835,10 +836,8 @@ def base_model(self): @pytest.fixture(autouse=True) def make_fleet_connector_not_abstract(self): - saved = FleetConnector.__abstractmethods__ - FleetConnector.__abstractmethods__ = frozenset() - yield - FleetConnector.__abstractmethods__ = saved + with _concrete(FleetConnector): + yield @pytest.fixture def fleet_connector(self, base_model, mock_robot_session_pool): @@ -852,12 +851,8 @@ def fleet_connector(self, base_model, mock_robot_session_pool): ], ) ) - # Simulate the state after __connect() -> update_fleet() created the sessions. - connector._FleetConnector__robot_sessions = { - robot_id: mock_robot_session_pool.get_session(robot_id) - for robot_id in connector.robot_ids - } - return connector + # Seed the state after __connect() created the sessions. + return _seed_sessions(connector) @staticmethod def _sessions(connector): @@ -913,14 +908,21 @@ def test_remove_robot_unknown_is_noop( def test_publish_to_removed_robot_does_not_resurrect_session( self, fleet_connector, mock_robot_session_pool ): - """Accessing/publishing to a removed robot must raise, not silently recreate a - zombie session via the pool.""" + """Accessing a removed robot returns None and publishing to it is a graceful + skip that must NOT silently recreate a zombie session.""" fleet_connector.remove_robot("TestRobot1") - with pytest.raises(KeyError): - fleet_connector._get_robot_session("TestRobot1") - with pytest.raises(KeyError): - fleet_connector.publish_robot_odometry("TestRobot1", linear_speed=1.0) + # The accessor returns None for a robot with no active session. + assert fleet_connector._get_robot_session("TestRobot1") is None + + # Publishing tolerates the missing robot: no raise, and no resurrection via the + # pool. + mock_robot_session_pool.get_session.reset_mock() + fleet_connector.publish_robot_odometry("TestRobot1", linear_speed=1.0) + fleet_connector.publish_robot_pose("TestRobot1", 1.0, 2.0, 3.0, "map1") + fleet_connector.publish_robot_key_values("TestRobot1", foo="bar") + assert "TestRobot1" not in self._sessions(fleet_connector) + mock_robot_session_pool.get_session.assert_not_called() def test_update_fleet_reconciles_diff( self, fleet_connector, mock_robot_session_pool @@ -1091,9 +1093,114 @@ def test_cameras_registered_once_at_startup( finally: connector.stop() - def test_get_robot_config(self, fleet_connector): - assert fleet_connector._get_robot_config("TestRobot1").robot_id == "TestRobot1" - assert fleet_connector._get_robot_config("nope") is None + def test_update_fleet_keeps_config_of_unchanged_robot( + self, fleet_connector, mock_robot_session_pool + ): + """A robot in both the old and new fleet keeps its existing RobotConfig and + session even when the incoming RobotConfig differs (e.g. its cameras changed), + so config.fleet stays consistent with the live session.""" + cam_a = CameraConfig(video_url="rtsp://example/a") + fleet_connector.add_robot(RobotConfig(robot_id="CamBot", cameras=[cam_a])) + session = mock_robot_session_pool.get_session("CamBot") + session.register_camera.reset_mock() + mock_robot_session_pool.get_session.reset_mock() + mock_robot_session_pool.free_robot_session.reset_mock() + + # Re-declare the full fleet with CamBot's camera changed. + fleet_connector.update_fleet( + [ + RobotConfig(robot_id="TestRobot1"), + RobotConfig(robot_id="TestRobot2"), + RobotConfig( + robot_id="CamBot", + cameras=[CameraConfig(video_url="rtsp://example/b")], + ), + ] + ) + + # config.fleet still reports the original config; the session was untouched. + cam_bot_config = next( + rc for rc in fleet_connector.config.fleet if rc.robot_id == "CamBot" + ) + assert cam_bot_config.cameras == [cam_a] + mock_robot_session_pool.free_robot_session.assert_not_called() + session.register_camera.assert_not_called() + + def test_publish_skips_removed_robot_without_disrupting_others( + self, fleet_connector, mock_robot_session_pool + ): + """Publishing to a concurrently-removed robot is a graceful no-op skip; a robot + still in the fleet keeps publishing in the same pass.""" + session2 = fleet_connector._get_robot_session("TestRobot2") + fleet_connector.remove_robot("TestRobot1") + + # None of these raise for the removed robot. + fleet_connector.publish_robot_pose("TestRobot1", 1.0, 2.0, 3.0, "map1") + fleet_connector.publish_robot_odometry("TestRobot1", linear_speed=1.0) + fleet_connector.publish_robot_key_values("TestRobot1", foo="bar") + fleet_connector.publish_robot_map("TestRobot1", "map1") + + # A still-present robot publishes normally. + fleet_connector.publish_robot_odometry("TestRobot2", linear_speed=2.0) + session2.publish_odometry.assert_called_once_with(linear_speed=2.0) + + def test_add_remove_bypass_overridden_update_fleet( + self, base_model, mock_robot_session_pool + ): + """add_robot/remove_robot reconcile via the private path, so a subclass that + overrides the public update_fleet cannot intercept or break them.""" + override_calls = [] + + class CustomFleetConnector(FleetConnector): + def update_fleet(self, fleet): + override_calls.append(list(fleet)) # no reconcile + + CustomFleetConnector.__abstractmethods__ = frozenset() + connector = _seed_sessions( + CustomFleetConnector( + ConnectorRootConfig(**base_model, fleet=[RobotConfig(robot_id="R1")]) + ) + ) + + connector.add_robot(RobotConfig(robot_id="R2")) + assert "R2" in connector.robot_ids + assert "R2" in self._sessions(connector) + + connector.remove_robot("R1") + assert "R1" not in connector.robot_ids + assert "R1" not in self._sessions(connector) + + # The overridden public update_fleet was never invoked by the wrappers. + assert override_calls == [] + + def test_connect_logs_initialized_session_count( + self, base_model, mock_robot_session_pool + ): + """__connect logs how many sessions it created so an empty/misconfigured fleet + is not silent.""" + # Capture directly off the connector logger: it inherits root's WARNING level + # and may not propagate to caplog under the project's logging setup. + logger = logging.getLogger("inorbit_connector.connector") + messages = [] + handler = logging.Handler() + handler.emit = lambda record: messages.append(record.getMessage()) + logger.addHandler(handler) + old_level = logger.level + logger.setLevel(logging.INFO) + try: + connector = FleetConnector( + ConnectorRootConfig(**base_model, fleet=[RobotConfig(robot_id="R1")]) + ) + connector.start() + try: + sleep(0.5) + finally: + connector.stop() + finally: + logger.removeHandler(handler) + logger.setLevel(old_level) + + assert any("Initialized 1 robot session(s)" in m for m in messages) def test_thread_safety_smoke(self, fleet_connector): """Concurrent add/remove on a SHARED id set must keep the fleet's internal @@ -1199,10 +1306,8 @@ def base_model(self): @pytest.fixture(autouse=True) def make_connector_not_abstract(self): - saved = Connector.__abstractmethods__ - Connector.__abstractmethods__ = frozenset() - yield - Connector.__abstractmethods__ = saved + with _concrete(Connector): + yield @pytest.fixture def base_connector(self, base_model): @@ -1250,6 +1355,18 @@ def test_get_session(self, base_connector, mock_robot_session_pool): assert session is not None assert session.robot_id == "TestRobot" + def test_get_session_raises_without_active_session( + self, base_model, mock_robot_session_pool + ): + """_get_session keeps the strict contract: it raises (not returns None) when + the robot has no active session, e.g. before the connector connects.""" + connector = Connector( + "TestRobot", + ConnectorRootConfig(**base_model, fleet=[RobotConfig(robot_id="TestRobot")]), + ) + with pytest.raises(KeyError): + connector._get_session() + def test_fleet_mutation_methods_raise(self, base_connector): """Single-robot Connector blocks the runtime fleet-mutation API.""" with pytest.raises(NotImplementedError): @@ -1538,10 +1655,8 @@ def base_model(self): @pytest.fixture(autouse=True) def make_connector_not_abstract(self): - saved = Connector.__abstractmethods__ - Connector.__abstractmethods__ = frozenset() - yield - Connector.__abstractmethods__ = saved + with _concrete(Connector): + yield @pytest.fixture def base_connector(self, base_model): From e1be4ae4d8d8f203d46907a6c52f2f21531f49fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Badenes?= Date: Fri, 5 Jun 2026 11:39:09 -0300 Subject: [PATCH 11/11] Update test_connector.py --- tests/test_connector.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_connector.py b/tests/test_connector.py index 0269462..a96d314 100644 --- a/tests/test_connector.py +++ b/tests/test_connector.py @@ -1362,7 +1362,9 @@ def test_get_session_raises_without_active_session( the robot has no active session, e.g. before the connector connects.""" connector = Connector( "TestRobot", - ConnectorRootConfig(**base_model, fleet=[RobotConfig(robot_id="TestRobot")]), + ConnectorRootConfig( + **base_model, fleet=[RobotConfig(robot_id="TestRobot")] + ), ) with pytest.raises(KeyError): connector._get_session()