diff --git a/qtm_rt/protocol.py b/qtm_rt/protocol.py index 8ac0f0d..cb908bf 100644 --- a/qtm_rt/protocol.py +++ b/qtm_rt/protocol.py @@ -80,8 +80,13 @@ async def await_event(self, event=None, timeout=None): if self.event_future is not None: raise Exception("Can't wait on multiple events!") - result = await asyncio.wait_for(self._wait_loop(event), timeout) - return result + try: + return await asyncio.wait_for(self._wait_loop(event), timeout) + finally: + # On timeout/cancellation the future is still queued and pending. + # Clear it so the next await_event isn't blocked by stale state. + # On success _on_event already cleared it; assignment is a no-op. + self.event_future = None def send_command( self, command, callback=True, command_type=QRTPacketType.PacketCommand diff --git a/test/qtmprotocol_test.py b/test/qtmprotocol_test.py index 8ba60af..7dd8363 100644 --- a/test/qtmprotocol_test.py +++ b/test/qtmprotocol_test.py @@ -32,6 +32,22 @@ async def test_await_any_event_timeout(qtmprotocol: QTMProtocol): await awaitable +@pytest.mark.asyncio +async def test_await_event_after_timeout(qtmprotocol: QTMProtocol): + # First call times out — previously this left event_future set, which made + # the next await_event raise "Can't wait on multiple events!" even though + # nothing was actually being awaited. + with pytest.raises(asyncio.TimeoutError): + await qtmprotocol.await_event(timeout=0.05) + + # Second call must succeed; event_future was cleared on the timeout path. + awaitable = qtmprotocol.await_event(timeout=1) + asyncio.get_running_loop().call_later( + 0, lambda: qtmprotocol._on_event(QRTEvent.EventConnected) + ) + assert await awaitable == QRTEvent.EventConnected + + @pytest.mark.asyncio async def test_await_any_event(qtmprotocol: QTMProtocol): awaitable = qtmprotocol.await_event(timeout=1)