diff --git a/test/conftest.py b/test/conftest.py index 047f657f1..e20bee7dc 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -161,9 +161,11 @@ async def receive(): """Return a websocket message.""" await asyncio.sleep(0) - message = messages.popleft() - if not messages: + try: + message = messages.popleft() + except IndexError: ws_client.closed = True + return WSMessage(WSMsgType.CLOSED, None, None) return message diff --git a/test/test_client.py b/test/test_client.py index 507e78a6e..abf4c6293 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -461,3 +461,18 @@ async def test_additional_user_agent_components(client_session, url): }, } ) + + +async def test_pop_future_none(client_session, url, driver_ready): + """Test when a future has been cleared from futures dict, popping still works.""" + client = Client(url, client_session) + await client.connect() + + assert client.connected + + asyncio.create_task(client.listen(driver_ready)) + + await driver_ready.wait() + + with pytest.raises(asyncio.CancelledError): + await client.async_send_command({"command": "some_command"}) diff --git a/test/test_main.py b/test/test_main.py index 7094f0286..76074b4e9 100644 --- a/test/test_main.py +++ b/test/test_main.py @@ -72,4 +72,4 @@ def test_connect(client_session, url, ws_client): main() assert sys_exit.value.code == 0 - assert ws_client.receive.call_count == 4 + assert ws_client.receive.call_count == 5 diff --git a/zwave_js_server/client.py b/zwave_js_server/client.py index 4e73fe534..a9fe8794f 100644 --- a/zwave_js_server/client.py +++ b/zwave_js_server/client.py @@ -115,7 +115,7 @@ async def async_send_command( try: return await future finally: - self._result_futures.pop(message_id) + self._result_futures.pop(message_id, None) async def async_send_command_no_wait( self, message: dict[str, Any], require_schema: int | None = None