Fix wait_for_future for cancelled futures (#88)

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

bdai_ros2_wrappers/bdai_ros2_wrappers/futures.py

@@ -23,5 +23,7 @@ def wait_for_future(future: Future, timeout_sec: Optional[float] = None, *, cont
     event = Event()
     context.on_shutdown(event.set)
     future.add_done_callback(lambda _: event.set())
+    if future.cancelled():
+        event.set()
     event.wait(timeout=timeout_sec)
     return future.done()

bdai_ros2_wrappers/test/test_futures.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2024 Boston Dynamics AI Institute Inc.  All rights reserved.
+
+from rclpy.task import Future
+
+from bdai_ros2_wrappers.futures import wait_for_future
+from bdai_ros2_wrappers.scope import ROSAwareScope
+
+
+def test_wait_for_cancelled_future(ros: ROSAwareScope) -> None:
+    """Asserts that waiting for a cancelled future does not hang indefinitely."""
+    future = Future()
+    future.cancel()
+
+    assert not wait_for_future(future, context=ros.context)
+    assert future.cancelled()