reduce

homeassistant/components/dhcp/websocket_api.py

@@ -25,59 +25,6 @@ def async_setup(hass: HomeAssistant) -> None:
     websocket_api.async_register_command(hass, ws_subscribe_discovery)
 
 
-class _DiscoverySubscription:
-    """Class to hold and manage the subscription data."""
-
-    def __init__(
-        self,
-        hass: HomeAssistant,
-        connection: websocket_api.ActiveConnection,
-        ws_msg_id: int,
-    ) -> None:
-        """Initialize the subscription data."""
-        self.hass = hass
-        self.ws_msg_id = ws_msg_id
-        self.connection = connection
-
-    @callback
-    def async_start(self) -> None:
-        """Start the subscription."""
-        connection = self.connection
-        connection.subscriptions[self.ws_msg_id] = (
-            async_register_dhcp_callback_internal(
-                self.hass,
-                self._async_send_address_data,
-            )
-        )
-        connection.send_message(
-            json_bytes(websocket_api.result_message(self.ws_msg_id))
-        )
-        self._async_send_address_data(
-            async_get_address_data_internal(self.hass),
-        )
-
-    def _async_event_message(self, message: dict[str, Any]) -> None:
-        self.connection.send_message(
-            json_bytes(websocket_api.event_message(self.ws_msg_id, message))
-        )
-
-    def _async_send_address_data(
-        self, address_data: dict[str, DHCPAddressData]
-    ) -> None:
-        self._async_event_message(
-            {
-                "add": [
-                    {
-                        "mac_address": dr.format_mac(mac_address).upper(),
-                        "hostname": data[HOSTNAME],
-                        "ip_address": data[IP_ADDRESS],
-                    }
-                    for mac_address, data in address_data.items()
-                ]
-            }
-        )
-
-
 @websocket_api.require_admin
 @websocket_api.websocket_command(
     {
@@ -89,4 +36,28 @@ async def ws_subscribe_discovery(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
 ) -> None:
     """Handle subscribe discovery websocket command."""
-    _DiscoverySubscription(hass, connection, msg["id"]).async_start()
+    ws_msg_id: int = msg["id"]
+
+    def _async_send(address_data: dict[str, DHCPAddressData]) -> None:
+        connection.send_message(
+            json_bytes(
+                websocket_api.event_message(
+                    ws_msg_id,
+                    {
+                        "add": [
+                            {
+                                "mac_address": dr.format_mac(mac_address).upper(),
+                                "hostname": data[HOSTNAME],
+                                "ip_address": data[IP_ADDRESS],
+                            }
+                            for mac_address, data in address_data.items()
+                        ]
+                    },
+                )
+            )
+        )
+
+    unsub = async_register_dhcp_callback_internal(hass, _async_send)
+    connection.subscriptions[ws_msg_id] = unsub
+    connection.send_message(json_bytes(websocket_api.result_message(ws_msg_id)))
+    _async_send(async_get_address_data_internal(hass))