From 83484c04cf06808fe4f53b553fd9d07bf7bd0d41 Mon Sep 17 00:00:00 2001 From: Shane Loretz Date: Thu, 9 Dec 2021 17:22:30 -0800 Subject: [PATCH] Add __enter__ and __exit__ to JumpHandle Signed-off-by: Shane Loretz --- rclpy/rclpy/clock.py | 6 ++++++ rclpy/test/test_clock.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/rclpy/rclpy/clock.py b/rclpy/rclpy/clock.py index 5157a69da..de2ada854 100644 --- a/rclpy/rclpy/clock.py +++ b/rclpy/rclpy/clock.py @@ -95,6 +95,12 @@ def unregister(self): self._clock.handle.remove_clock_callback(self) self._clock = None + def __enter__(self): + return self + + def __exit__(self, t, v, tb): + self.unregister() + class Clock: diff --git a/rclpy/test/test_clock.py b/rclpy/test/test_clock.py index 9a7526f23..6b6246c1f 100644 --- a/rclpy/test/test_clock.py +++ b/rclpy/test/test_clock.py @@ -18,6 +18,7 @@ from rclpy.clock import Clock from rclpy.clock import ClockType +from rclpy.clock import JumpHandle from rclpy.clock import JumpThreshold from rclpy.clock import ROSClock from rclpy.duration import Duration @@ -179,3 +180,20 @@ def test_triggered_clock_change_callbacks(self): handler1.unregister() handler2.unregister() handler3.unregister() + + +def test_with_jump_handle(): + clock = ROSClock() + clock._set_ros_time_is_active(False) + + post_callback = Mock() + threshold = JumpThreshold(min_forward=None, min_backward=None, on_clock_change=True) + + with clock.create_jump_callback(threshold, post_callback=post_callback) as jump_handler: + assert isinstance(jump_handler, JumpHandle) + clock._set_ros_time_is_active(True) + post_callback.assert_called_once() + + post_callback.reset_mock() + clock._set_ros_time_is_active(False) + post_callback.assert_not_called()