From a1273c73e461ba2392d7d63db11f40cc7738aac6 Mon Sep 17 00:00:00 2001
From: mhidalgo-bdai <144129882+mhidalgo-bdai@users.noreply.github.com>
Date: Wed, 18 Oct 2023 15:29:42 -0300
Subject: [PATCH] Add `wait_for_shutdown` utility (#29)

Signed-off-by: Michel Hidalgo <mhidalgo@theaiinstitute.com>
---
 .../bdai_ros2_wrappers/context.py             | 24 +++++++++++++++++++
 .../bdai_ros2_wrappers/process.py             |  8 ++-----
 bdai_ros2_wrappers/test/test_context.py       | 14 +++++++++++
 3 files changed, 40 insertions(+), 6 deletions(-)
 create mode 100644 bdai_ros2_wrappers/bdai_ros2_wrappers/context.py
 create mode 100644 bdai_ros2_wrappers/test/test_context.py

diff --git a/bdai_ros2_wrappers/bdai_ros2_wrappers/context.py b/bdai_ros2_wrappers/bdai_ros2_wrappers/context.py
new file mode 100644
index 0000000..b47a04e
--- /dev/null
+++ b/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)
diff --git a/bdai_ros2_wrappers/bdai_ros2_wrappers/process.py b/bdai_ros2_wrappers/bdai_ros2_wrappers/process.py
index 3ffeef3..112361c 100644
--- a/bdai_ros2_wrappers/bdai_ros2_wrappers/process.py
+++ b/bdai_ros2_wrappers/bdai_ros2_wrappers/process.py
@@ -12,6 +12,7 @@
 import rclpy.logging
 import rclpy.node
 
+import bdai_ros2_wrappers.context as context
 import bdai_ros2_wrappers.scope as scope
 from bdai_ros2_wrappers.scope import AnyEntity, AnyEntityFactoryCallable, ROSAwareScope
 from bdai_ros2_wrappers.utilities import either_or
@@ -172,12 +173,7 @@ def wait_for_shutdown(self, *, timeout_sec: typing.Optional[float] = None) -> bo
         with self._lock:
             if self._scope is None:
                 raise RuntimeError("process is not executing")
-            context = self._scope.context
-            if context is None:
-                context = rclpy.get_default_context()
-            event = threading.Event()
-            context.on_shutdown(event.set)
-            return event.wait(timeout_sec)
+            return context.wait_for_shutdown(timeout_sec=timeout_sec, context=self._scope.context)
 
     def try_shutdown(self) -> None:
         """Atempts to shutdown the underlying scope context."""
diff --git a/bdai_ros2_wrappers/test/test_context.py b/bdai_ros2_wrappers/test/test_context.py
new file mode 100644
index 0000000..05a35ad
--- /dev/null
+++ b/bdai_ros2_wrappers/test/test_context.py
@@ -0,0 +1,14 @@
+# Copyright (c) 2023 Boston Dynamics AI Institute Inc.  All rights reserved.
+
+import rclpy
+
+from bdai_ros2_wrappers.context import wait_for_shutdown
+from bdai_ros2_wrappers.process import ROSAwareScope
+
+
+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)
+    assert ros.executor is not None
+    ros.executor.create_task(lambda: rclpy.shutdown())
+    assert wait_for_shutdown(timeout_sec=10.0)