Change execute_service to be async to handle timeouts internally (#1443)

Author: jesserockz

aioesphomeapi/client.py

@@ -169,6 +169,7 @@
 
 DEFAULT_BLE_TIMEOUT = 30.0
 DEFAULT_BLE_DISCONNECT_TIMEOUT = 20.0
+DEFAULT_EXECUTE_SERVICE_TIMEOUT = 30.0
 
 SUBSCRIBE_STATES_MSG_TYPES = (*SUBSCRIBE_STATES_RESPONSE_TYPES, CameraImageResponse)
 
@@ -1313,17 +1314,17 @@ def update_command(
             UpdateCommandRequest(key=key, command=command, device_id=device_id)
         )
 
-    def execute_service(
+    async def execute_service(
         self,
         service: UserService,
         data: ExecuteServiceDataType,
         *,
-        on_response: Callable[[ExecuteServiceResponseModel], None] | None = None,
-        return_response: bool = False,
-    ) -> None:
+        return_response: bool | None = None,
+        timeout: float = DEFAULT_EXECUTE_SERVICE_TIMEOUT,
+    ) -> ExecuteServiceResponseModel | None:
         connection = self._get_connection()
         # Generate call_id when response callback is provided
-        call_id = next(self._call_id_counter) if on_response is not None else 0
+        call_id = next(self._call_id_counter) if return_response is not None else 0
         req = ExecuteServiceRequest(
             key=service.key,
             call_id=call_id,
@@ -1353,21 +1354,31 @@ def execute_service(
         req.args.extend(args)
 
         # Register callback for response if provided
-        if on_response is not None:
-            unsub: Callable[[], None] | None = None
+        if return_response is not None:
+            response_event = asyncio.Event()
+            response_msg: ExecuteServiceResponseModel | None = None
 
             def _on_response(msg: ExecuteServiceResponse) -> None:
+                nonlocal response_msg
                 if msg.call_id == call_id:
-                    on_response(ExecuteServiceResponseModel.from_pb(msg))
-                    if unsub is not None:
-                        unsub()
+                    response_msg = ExecuteServiceResponseModel.from_pb(msg)
+                    response_event.set()
 
             unsub = connection.add_message_callback(
                 _on_response,
                 (ExecuteServiceResponse,),
             )
 
-        connection.send_message(req)
+            try:
+                connection.send_message(req)
+                await asyncio.wait_for(response_event.wait(), timeout=timeout)
+                return response_msg
+            finally:
+                unsub()
+        else:
+            connection.send_message(req)
+
+        return None
 
     def _request_image(self, *, single: bool = False, stream: bool = False) -> None:
         self._get_connection().send_message(

tests/test_client.py

@@ -122,7 +122,6 @@
     ClimateSwingMode,
     DeviceInfo,
     ESPHomeBluetoothGATTServices,
-    ExecuteServiceResponse as ExecuteServiceResponseModel,
     FanDirection,
     FanSpeed,
     HomeassistantActionResponse as HomeassistantActionResponseModel,
@@ -963,9 +962,9 @@ async def test_execute_service(auth_client: APIClient) -> None:
     )
 
     with pytest.raises(KeyError):
-        auth_client.execute_service(service, data={})
+        await auth_client.execute_service(service, data={})
 
-    auth_client.execute_service(
+    await auth_client.execute_service(
         service,
         data={
             "arg1": False,
@@ -1006,7 +1005,7 @@ async def test_execute_service(auth_client: APIClient) -> None:
     )
 
     # Test legacy_int
-    auth_client.execute_service(
+    await auth_client.execute_service(
         service,
         data={
             "arg1": False,
@@ -1025,7 +1024,7 @@ async def test_execute_service(auth_client: APIClient) -> None:
     send.reset_mock()
 
     # Test arg order
-    auth_client.execute_service(
+    await auth_client.execute_service(
         service,
         data={
             "arg2": 42,
@@ -1045,7 +1044,7 @@ async def test_execute_service(auth_client: APIClient) -> None:
 
 
 async def test_execute_service_with_call_id(auth_client: APIClient) -> None:
-    """Test that call_id is auto-generated when on_response is provided."""
+    """Test that call_id is auto-generated when return_response is set."""
     send = patch_send(auth_client)
     patch_api_version(auth_client, APIVersion(1, 3))
 
@@ -1057,51 +1056,45 @@ async def test_execute_service_with_call_id(auth_client: APIClient) -> None:
         ],
     )
 
-    def dummy_callback(response: ExecuteServiceResponseModel) -> None:
-        pass
-
-    # Test without on_response - call_id should be 0
-    auth_client.execute_service(
+    # Test without return_response - call_id should be 0
+    await auth_client.execute_service(
         service,
         data={"arg1": True},
     )
-    send.assert_called_once_with(
-        ExecuteServiceRequest(
-            key=1,
-            args=[
-                ExecuteServiceArgument(bool_=True),
-            ],
-            call_id=0,
-            return_response=False,
-        )
-    )
+    req = send.call_args[0][0]
+    assert req.call_id == 0
     send.reset_mock()
 
-    # Test with on_response - call_id should be auto-generated (non-zero)
-    auth_client.execute_service(
-        service,
-        data={"arg1": False},
-        on_response=dummy_callback,
-    )
+    # Test with return_response=True - call_id should be auto-generated (non-zero)
+    # Use short timeout since no response will come, we just want to verify the request
+    with pytest.raises(asyncio.TimeoutError):
+        await auth_client.execute_service(
+            service,
+            data={"arg1": False},
+            return_response=True,
+            timeout=0.01,
+        )
     req = send.call_args[0][0]
     assert req.call_id != 0  # Auto-generated
     first_call_id = req.call_id
     send.reset_mock()
 
     # Test that call_id increments
-    auth_client.execute_service(
-        service,
-        data={"arg1": True},
-        on_response=dummy_callback,
-    )
+    with pytest.raises(asyncio.TimeoutError):
+        await auth_client.execute_service(
+            service,
+            data={"arg1": True},
+            return_response=True,
+            timeout=0.01,
+        )
     req = send.call_args[0][0]
     assert req.call_id == first_call_id + 1
 
 
 async def test_execute_service_return_response_combinations(
     auth_client: APIClient,
 ) -> None:
-    """Test that return_response is passed through and call_id is generated for on_response."""
+    """Test return_response behavior and call_id generation."""
     send = patch_send(auth_client)
     patch_api_version(auth_client, APIVersion(1, 3))
 
@@ -1111,32 +1104,26 @@ async def test_execute_service_return_response_combinations(
         args=[],
     )
 
-    def dummy_callback(response: ExecuteServiceResponseModel) -> None:
-        pass
-
-    # Case 1: no callback, no return_response -> call_id=0, return_response=False
-    auth_client.execute_service(service, data={})
-    assert send.call_args[0][0].return_response is False
+    # Case 1: return_response=None (default) -> call_id=0, no waiting
+    await auth_client.execute_service(service, data={})
     assert send.call_args[0][0].call_id == 0
     send.reset_mock()
 
-    # Case 2: return_response=True without callback -> call_id=0, return_response=True
-    auth_client.execute_service(service, data={}, return_response=True)
+    # Case 2: return_response=True -> generates call_id, waits for response
+    with pytest.raises(asyncio.TimeoutError):
+        await auth_client.execute_service(
+            service, data={}, return_response=True, timeout=0.01
+        )
     assert send.call_args[0][0].return_response is True
-    assert send.call_args[0][0].call_id == 0
-    send.reset_mock()
-
-    # Case 3: on_response generates call_id, return_response passed through
-    auth_client.execute_service(service, data={}, on_response=dummy_callback)
-    assert send.call_args[0][0].return_response is False
     assert send.call_args[0][0].call_id != 0
     send.reset_mock()
 
-    # Case 4: on_response with return_response=True
-    auth_client.execute_service(
-        service, data={}, on_response=dummy_callback, return_response=True
-    )
-    assert send.call_args[0][0].return_response is True
+    # Case 3: return_response=False -> generates call_id, waits for response
+    with pytest.raises(asyncio.TimeoutError):
+        await auth_client.execute_service(
+            service, data={}, return_response=False, timeout=0.01
+        )
+    assert send.call_args[0][0].return_response is False
     assert send.call_args[0][0].call_id != 0
     send.reset_mock()
 
@@ -2249,15 +2236,14 @@ def on_zwave_proxy_request(msg: ZWaveProxyRequest) -> None:
     assert first_msg.data == b"\x00\x01\x02\x03"
 
 
-async def test_execute_service_with_response_callback(
+async def test_execute_service_with_response(
     api_client: tuple[
         APIClient, APIConnection, asyncio.Transport, APIPlaintextFrameHelper
     ],
 ) -> None:
-    """Test execute_service with on_response callback."""
+    """Test execute_service with return_response returns response directly."""
     client, connection, _transport, protocol = api_client
     patch_api_version(client, APIVersion(1, 3))
-    test_msg: list[ExecuteServiceResponseModel] = []
     sent_requests: list[ExecuteServiceRequest] = []
 
     # Capture sent requests to get auto-generated call_id
@@ -2270,9 +2256,6 @@ def capture_send(msg: Any) -> None:
 
     connection.send_message = capture_send
 
-    def on_response(msg: ExecuteServiceResponseModel) -> None:
-        test_msg.append(msg)
-
     service = UserService(
         name="my_service",
         key=1,
@@ -2281,13 +2264,15 @@ def on_response(msg: ExecuteServiceResponseModel) -> None:
         ],
     )
 
-    # Execute service with callback
-    client.execute_service(
-        service,
-        data={"arg1": True},
-        on_response=on_response,
+    # Execute service with return_response - start as task so we can simulate response
+    task = asyncio.create_task(
+        client.execute_service(
+            service,
+            data={"arg1": True},
+            return_response=True,
+        )
     )
-    await asyncio.sleep(0)
+    await asyncio.sleep(0)  # Let task start and send request
 
     # Get the auto-generated call_id
     assert len(sent_requests) == 1
@@ -2302,31 +2287,22 @@ def on_response(msg: ExecuteServiceResponseModel) -> None:
         response_data=b'{"result": "ok"}',
     )
     mock_data_received(protocol, generate_plaintext_packet(response))
+    result = await task  # Task should complete now that response was received
 
-    assert len(test_msg) == 1
-    first_msg = test_msg[0]
-    assert first_msg.call_id == first_call_id
-    assert first_msg.success is True
-    assert first_msg.error_message == ""
-    assert first_msg.response_data == b'{"result": "ok"}'
+    assert result is not None
+    assert result.call_id == first_call_id
+    assert result.success is True
+    assert result.error_message == ""
+    assert result.response_data == b'{"result": "ok"}'
 
-    # Callback should auto-unsubscribe, so another response shouldn't be received
-    response2: message.Message = ExecuteServiceResponsePb(
-        call_id=first_call_id,
-        success=True,
-        error_message="",
-        response_data=b'{"result": "second"}',
-    )
-    mock_data_received(protocol, generate_plaintext_packet(response2))
-    assert len(test_msg) == 1  # Still only one message
-
-    # Test that responses with different call_id are ignored
-    test_msg.clear()
+    # Test that responses with different call_id are ignored until correct one arrives
     sent_requests.clear()
-    client.execute_service(
-        service,
-        data={"arg1": False},
-        on_response=on_response,
+    task2 = asyncio.create_task(
+        client.execute_service(
+            service,
+            data={"arg1": False},
+            return_response=True,
+        )
     )
     await asyncio.sleep(0)
 
@@ -2343,7 +2319,8 @@ def on_response(msg: ExecuteServiceResponseModel) -> None:
         response_data=b"",
     )
     mock_data_received(protocol, generate_plaintext_packet(wrong_response))
-    assert len(test_msg) == 0
+    await asyncio.sleep(0)
+    assert not task2.done()  # Task still waiting
 
     # Correct call_id should be received
     correct_response: message.Message = ExecuteServiceResponsePb(
@@ -2353,8 +2330,9 @@ def on_response(msg: ExecuteServiceResponseModel) -> None:
         response_data=b"",
     )
     mock_data_received(protocol, generate_plaintext_packet(correct_response))
-    assert len(test_msg) == 1
-    assert test_msg[0].call_id == second_call_id
+    result2 = await task2  # Task should complete now
+    assert result2 is not None
+    assert result2.call_id == second_call_id
 
 
 async def test_subscribe_service_calls(auth_client: APIClient) -> None:
@@ -3368,7 +3346,7 @@ async def test_calls_after_connection_closed(
         args=[],
     )
     with pytest.raises(APIConnectionError):
-        client.execute_service(service, {})
+        await client.execute_service(service, {})
     for method in (
         client.button_command,
         client.climate_command,