Add wait_for_shutdown utility

Signed-off-by: Michel Hidalgo <[email protected]>

bdai_ros2_wrappers/bdai_ros2_wrappers/context.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2023 Boston Dynamics AI Institute Inc.  All rights reserved.
+import threading
+from typing import Optional
+
+from rclpy.context import Context
+from rclpy.utilities import get_default_context
+
+
+def wait_for_shutdown(*, timeout_sec: Optional[float] = None, context: Optional[Context] = None) -> bool:
+    """
+    Wait for context shutdown.
+
+    Args:
+        timeout_sec: optional timeout for wait, wait indefinitely by default.
+        context: context to wait on, use default context by default.
+
+    Returns:
+        True if shutdown, False on timeout.
+    """
+    if context is None:
+        context = get_default_context()
+    event = threading.Event()
+    context.on_shutdown(event.set)
+    return event.wait(timeout_sec)

bdai_ros2_wrappers/test/test_context.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2023 Boston Dynamics AI Institute Inc.  All rights reserved.
+
+from typing import Iterable
+
+import pytest
+import rclpy
+
+from bdai_ros2_wrappers.context import wait_for_shutdown
+from bdai_ros2_wrappers.process import ROSAwareScope
+
+
+@pytest.fixture
+def ros() -> Iterable[ROSAwareScope]:
+    rclpy.init()
+    try:
+        with ROSAwareScope("fixture") as scope:
+            yield scope
+    finally:
+        rclpy.try_shutdown()
+
+
+def test_wait_for_shutdown(ros: ROSAwareScope) -> None:
+    """Asserts that wait for shutdown works as expected."""
+    assert not wait_for_shutdown(timeout_sec=1.0)
+    ros.executor.create_task(lambda: rclpy.shutdown())
+    assert wait_for_shutdown(timeout_sec=10.0)