From e23d35c6f067e898244905c4912ae6e0a92409e7 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Tue, 12 Oct 2021 11:32:57 +0200
Subject: [PATCH] Move all Tuya device handling into device listener class
 (#57523)

---
 homeassistant/components/tuya/__init__.py | 33 +++++++++++------------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/tuya/__init__.py b/homeassistant/components/tuya/__init__.py
index 326cacfebbf..8dd9979f3e7 100644
--- a/homeassistant/components/tuya/__init__.py
+++ b/homeassistant/components/tuya/__init__.py
@@ -132,20 +132,6 @@ async def cleanup_device_registry(
                 break
 
 
-@callback
-def async_remove_hass_device(
-    hass: HomeAssistant, entry: ConfigEntry, device_id: str
-) -> None:
-    """Remove device from hass cache."""
-    device_registry_object = device_registry.async_get(hass)
-    device_entry = device_registry_object.async_get_device(
-        identifiers={(DOMAIN, device_id)}
-    )
-    if device_entry is not None:
-        device_registry_object.async_remove_device(device_entry.id)
-        hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].discard(device_id)
-
-
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unloading the Tuya platforms."""
     _LOGGER.debug("integration unload")
@@ -191,9 +177,7 @@ class DeviceListener(TuyaDeviceListener):
             *self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_TUYA_MAP].values()
         ):
             ha_tuya_map = self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_TUYA_MAP]
-            self.hass.add_job(
-                async_remove_hass_device, self.hass, self.entry, device.id
-            )
+            self.hass.add_job(self.async_remove_device, device.id)
 
             for domain, tuya_list in ha_tuya_map.items():
                 if device.category in tuya_list:
@@ -224,4 +208,17 @@ class DeviceListener(TuyaDeviceListener):
     def remove_device(self, device_id: str) -> None:
         """Add device removed listener."""
         _LOGGER.debug("tuya remove device:%s", device_id)
-        self.hass.add_job(async_remove_hass_device, self.hass, self.entry, device_id)
+        self.hass.add_job(self.async_remove_device, device_id)
+
+    @callback
+    def async_remove_device(self, device_id: str) -> None:
+        """Remove device from Home Assistant."""
+        device_registry_object = device_registry.async_get(self.hass)
+        device_entry = device_registry_object.async_get_device(
+            identifiers={(DOMAIN, device_id)}
+        )
+        if device_entry is not None:
+            device_registry_object.async_remove_device(device_entry.id)
+            self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_DEVICES].discard(
+                device_id
+            )