From 6d809b0b5ad5993a973c21122f4da713645b96e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ab=C3=ADlio=20Costa?= Date: Tue, 13 May 2025 21:57:15 +0100 Subject: [PATCH] Add service response support to admin services (#144837) --- homeassistant/helpers/service.py | 29 ++++++++++++++++++++++------- tests/helpers/test_service.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/homeassistant/helpers/service.py b/homeassistant/helpers/service.py index 4873d935537..f157e82bc53 100644 --- a/homeassistant/helpers/service.py +++ b/homeassistant/helpers/service.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from collections.abc import Awaitable, Callable, Coroutine, Iterable +from collections.abc import Callable, Coroutine, Iterable import dataclasses from enum import Enum from functools import cache, partial @@ -1094,9 +1094,15 @@ async def _handle_entity_call( async def _async_admin_handler( hass: HomeAssistant, - service_job: HassJob[[ServiceCall], Awaitable[None] | None], + service_job: HassJob[ + [ServiceCall], + Coroutine[Any, Any, ServiceResponse | EntityServiceResponse] + | ServiceResponse + | EntityServiceResponse + | None, + ], call: ServiceCall, -) -> None: +) -> ServiceResponse | EntityServiceResponse | None: """Run an admin service.""" if call.context.user_id: user = await hass.auth.async_get_user(call.context.user_id) @@ -1105,9 +1111,10 @@ async def _async_admin_handler( if not user.is_admin: raise Unauthorized(context=call.context) - result = hass.async_run_hass_job(service_job, call) - if result is not None: - await result + task = hass.async_run_hass_job(service_job, call) + if task is not None: + return await task + return None @bind_hass @@ -1116,8 +1123,15 @@ def async_register_admin_service( hass: HomeAssistant, domain: str, service: str, - service_func: Callable[[ServiceCall], Awaitable[None] | None], + service_func: Callable[ + [ServiceCall], + Coroutine[Any, Any, ServiceResponse | EntityServiceResponse] + | ServiceResponse + | EntityServiceResponse + | None, + ], schema: VolSchemaType = vol.Schema({}, extra=vol.PREVENT_EXTRA), + supports_response: SupportsResponse = SupportsResponse.NONE, ) -> None: """Register a service that requires admin access.""" hass.services.async_register( @@ -1129,6 +1143,7 @@ def async_register_admin_service( HassJob(service_func, f"admin service {domain}.{service}"), ), schema, + supports_response, ) diff --git a/tests/helpers/test_service.py b/tests/helpers/test_service.py index 4582bce3e05..38e7e1ae452 100644 --- a/tests/helpers/test_service.py +++ b/tests/helpers/test_service.py @@ -32,6 +32,7 @@ from homeassistant.core import ( HassJob, HomeAssistant, ServiceCall, + ServiceResponse, SupportsResponse, ) from homeassistant.helpers import ( @@ -1648,6 +1649,33 @@ async def test_register_admin_service( assert calls[0].context.user_id == hass_admin_user.id +@pytest.mark.parametrize( + "supports_response", + [SupportsResponse.ONLY, SupportsResponse.OPTIONAL], +) +async def test_register_admin_service_return_response( + hass: HomeAssistant, supports_response: SupportsResponse +) -> None: + """Test the register admin service for a service that returns response data.""" + + async def mock_service(call: ServiceCall) -> ServiceResponse: + """Service handler coroutine.""" + assert call.return_response + return {"test-reply": "test-value1"} + + service.async_register_admin_service( + hass, "test", "test", mock_service, supports_response=supports_response + ) + result = await hass.services.async_call( + "test", + "test", + service_data={}, + blocking=True, + return_response=True, + ) + assert result == {"test-reply": "test-value1"} + + async def test_domain_control_not_async(hass: HomeAssistant, mock_entities) -> None: """Test domain verification in a service call with an unknown user.""" calls = []