Async response all the things (#17073)

* Use async_response

* Update device_registry.py

homeassistant/components/config/auth.py

@@ -1,9 +1,7 @@
 """Offer API to configure Home Assistant auth."""
 import voluptuous as vol
 
-from homeassistant.core import callback
 from homeassistant.components import websocket_api
-from homeassistant.components.websocket_api.decorators import require_owner
 
 
 WS_TYPE_LIST = 'config/auth/list'
@@ -41,26 +39,20 @@ async def async_setup(hass):
     return True
 
 
-@callback
-@require_owner
-def websocket_list(hass, connection, msg):
+@websocket_api.require_owner
+@websocket_api.async_response
+async def websocket_list(hass, connection, msg):
     """Return a list of users."""
-    async def send_users():
-        """Send users."""
     result = [_user_info(u) for u in await hass.auth.async_get_users()]
 
     connection.send_message(
         websocket_api.result_message(msg['id'], result))
 
-    hass.async_create_task(send_users())
 
-
-@callback
-@require_owner
-def websocket_delete(hass, connection, msg):
+@websocket_api.require_owner
+@websocket_api.async_response
+async def websocket_delete(hass, connection, msg):
     """Delete a user."""
-    async def delete_user():
-        """Delete user."""
     if msg['user_id'] == connection.user.id:
         connection.send_message(websocket_api.error_message(
             msg['id'], 'no_delete_self',
@@ -79,14 +71,10 @@ def websocket_delete(hass, connection, msg):
     connection.send_message(
         websocket_api.result_message(msg['id']))
 
-    hass.async_create_task(delete_user())
 
-
-@callback
-@require_owner
-def websocket_create(hass, connection, msg):
-    """Create a user."""
-    async def create_user():
+@websocket_api.require_owner
+@websocket_api.async_response
+async def websocket_create(hass, connection, msg):
     """Create a user."""
     user = await hass.auth.async_create_user(msg['name'])
 
@@ -95,8 +83,6 @@ def websocket_create(hass, connection, msg):
             'user': _user_info(user)
         }))
 
-    hass.async_create_task(create_user())
-
 
 def _user_info(user):
     """Format a user."""

homeassistant/components/config/device_registry.py

@@ -1,7 +1,6 @@
 """HTTP views to interact with the device registry."""
 import voluptuous as vol
 
-from homeassistant.core import callback
 from homeassistant.helpers.device_registry import async_get_registry
 from homeassistant.components import websocket_api
 
@@ -22,14 +21,9 @@ async def async_setup(hass):
     return True
 
 
-@callback
-def websocket_list_devices(hass, connection, msg):
-    """Handle list devices command.
-
-    Async friendly.
-    """
-    async def retrieve_entities():
-        """Get devices from registry."""
+@websocket_api.async_response
+async def websocket_list_devices(hass, connection, msg):
+    """Handle list devices command."""
     registry = await async_get_registry(hass)
     connection.send_message(websocket_api.result_message(
         msg['id'], [{
@@ -43,5 +37,3 @@ def websocket_list_devices(hass, connection, msg):
             'hub_device_id': entry.hub_device_id,
         } for entry in registry.devices.values()]
     ))
-
-    hass.async_create_task(retrieve_entities())

homeassistant/components/websocket_api/__init__.py

@@ -20,6 +20,7 @@ BASE_COMMAND_MESSAGE_SCHEMA = messages.BASE_COMMAND_MESSAGE_SCHEMA
 error_message = messages.error_message
 result_message = messages.result_message
 async_response = decorators.async_response
+require_owner = decorators.require_owner
 ws_require_user = decorators.ws_require_user
 # pylint: enable=invalid-name
 

homeassistant/components/websocket_api/decorators.py

@@ -10,9 +10,7 @@ from . import messages
 _LOGGER = logging.getLogger(__name__)
 
 
-def async_response(func):
-    """Decorate an async function to handle WebSocket API messages."""
-    async def handle_msg_response(hass, connection, msg):
+async def _handle_async_response(func, hass, connection, msg):
     """Create a response and handle exception."""
     try:
         await func(hass, connection, msg)
@@ -21,11 +19,15 @@ def async_response(func):
         connection.send_message(messages.error_message(
             msg['id'], 'unknown', 'Unexpected error occurred'))
 
+
+def async_response(func):
+    """Decorate an async function to handle WebSocket API messages."""
     @callback
     @wraps(func)
     def schedule_handler(hass, connection, msg):
         """Schedule the handler."""
-        hass.async_create_task(handle_msg_response(hass, connection, msg))
+        hass.async_create_task(
+            _handle_async_response(func, hass, connection, msg))
 
     return schedule_handler