From 679b2f4a292e9b8135d7dc1ff098afd100229c30 Mon Sep 17 00:00:00 2001
From: mhidalgo-bdai <144129882+mhidalgo-bdai@users.noreply.github.com>
Date: Thu, 11 Apr 2024 15:05:20 -0300
Subject: [PATCH] Fix `wait_for_future` for cancelled futures (#88)

Signed-off-by: Michel Hidalgo <mhidalgo@theaiinstitute.com>
---
 bdai_ros2_wrappers/bdai_ros2_wrappers/futures.py |  2 ++
 bdai_ros2_wrappers/test/test_futures.py          | 15 +++++++++++++++
 2 files changed, 17 insertions(+)
 create mode 100644 bdai_ros2_wrappers/test/test_futures.py

diff --git a/bdai_ros2_wrappers/bdai_ros2_wrappers/futures.py b/bdai_ros2_wrappers/bdai_ros2_wrappers/futures.py
index 02c7805..c48577a 100644
--- a/bdai_ros2_wrappers/bdai_ros2_wrappers/futures.py
+++ b/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()
diff --git a/bdai_ros2_wrappers/test/test_futures.py b/bdai_ros2_wrappers/test/test_futures.py
new file mode 100644
index 0000000..129a0e9
--- /dev/null
+++ b/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()