From b23a66d776c41cbe484a38ab94dd7db588c3bde6 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Thu, 20 Oct 2022 12:20:39 +0200
Subject: [PATCH] Add websocket type hints in entity_registry (#80657)

* Add websocket type hints in entity_registry

* Adjust websocket_list_entities

* Fix update

* Fix websocket_update_entity

* Apply suggestion

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* Apply suggestion

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
---
 .../components/config/entity_registry.py      | 37 ++++++++++++++-----
 .../components/config/test_entity_registry.py |  2 +-
 2 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/config/entity_registry.py b/homeassistant/components/config/entity_registry.py
index b024c3f0128..b4bd7403c43 100644
--- a/homeassistant/components/config/entity_registry.py
+++ b/homeassistant/components/config/entity_registry.py
@@ -36,14 +36,18 @@ async def async_setup(hass: HomeAssistant) -> bool:
         {vol.Required("type"): "config/entity_registry/list"}
     )
     @callback
-    def websocket_list_entities(hass, connection, msg):
+    def websocket_list_entities(
+        hass: HomeAssistant,
+        connection: websocket_api.ActiveConnection,
+        msg: dict[str, Any],
+    ) -> None:
         """Handle list registry entries command."""
         nonlocal cached_list_entities
         if not cached_list_entities:
             registry = er.async_get(hass)
             cached_list_entities = message_to_json(
                 websocket_api.result_message(
-                    IDEN_TEMPLATE,
+                    IDEN_TEMPLATE,  # type: ignore[arg-type]
                     [_entry_dict(entry) for entry in registry.entities.values()],
                 )
             )
@@ -70,7 +74,11 @@ async def async_setup(hass: HomeAssistant) -> bool:
     }
 )
 @callback
-def websocket_get_entity(hass, connection, msg):
+def websocket_get_entity(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict[str, Any],
+) -> None:
     """Handle get entity registry entry command.
 
     Async friendly.
@@ -120,7 +128,11 @@ def websocket_get_entity(hass, connection, msg):
     }
 )
 @callback
-def websocket_update_entity(hass, connection, msg):
+def websocket_update_entity(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict[str, Any],
+) -> None:
     """Handle update entity websocket command.
 
     Async friendly.
@@ -153,7 +165,7 @@ def websocket_update_entity(hass, connection, msg):
         if entity_entry.device_id:
             device_registry = dr.async_get(hass)
             device = device_registry.async_get(entity_entry.device_id)
-            if device.disabled:
+            if device and device.disabled:
                 connection.send_message(
                     websocket_api.error_message(
                         msg["id"], "invalid_info", "Device is disabled"
@@ -184,11 +196,14 @@ def websocket_update_entity(hass, connection, msg):
         )
         return
 
-    result = {"entity_entry": _entry_ext_dict(entity_entry)}
+    result: dict[str, Any] = {"entity_entry": _entry_ext_dict(entity_entry)}
     if "disabled_by" in changes and changes["disabled_by"] is None:
         # Enabling an entity requires a config entry reload, or HA restart
-        config_entry = hass.config_entries.async_get_entry(entity_entry.config_entry_id)
-        if config_entry and not config_entry.supports_unload:
+        if (
+            not (config_entry_id := entity_entry.config_entry_id)
+            or (config_entry := hass.config_entries.async_get_entry(config_entry_id))
+            and not config_entry.supports_unload
+        ):
             result["require_restart"] = True
         else:
             result["reload_delay"] = config_entries.RELOAD_AFTER_UPDATE_DELAY
@@ -203,7 +218,11 @@ def websocket_update_entity(hass, connection, msg):
     }
 )
 @callback
-def websocket_remove_entity(hass, connection, msg):
+def websocket_remove_entity(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict[str, Any],
+) -> None:
     """Handle remove entity websocket command.
 
     Async friendly.
diff --git a/tests/components/config/test_entity_registry.py b/tests/components/config/test_entity_registry.py
index 2f4cd980d8e..ee46838b01c 100644
--- a/tests/components/config/test_entity_registry.py
+++ b/tests/components/config/test_entity_registry.py
@@ -345,7 +345,7 @@ async def test_update_entity(hass, client):
             "platform": "test_platform",
             "unique_id": "1234",
         },
-        "reload_delay": 30,
+        "require_restart": True,
     }
 
     # UPDATE ENTITY OPTION