Refactor Tuya device handling (#57545)

* Refactor Tuya device handling

* Tweak

homeassistant/components/tuya/__init__.py

@@ -1,7 +1,8 @@
 """Support for Tuya Smart devices."""
+from __future__ import annotations
 
-import itertools
 import logging
+from typing import NamedTuple
 
 from tuya_iot import (
     AuthType,
@@ -30,24 +31,25 @@ from .const import (
     CONF_USERNAME,
     DOMAIN,
     PLATFORMS,
-    TUYA_DEVICE_MANAGER,
     TUYA_DISCOVERY_NEW,
-    TUYA_HA_DEVICES,
     TUYA_HA_SIGNAL_UPDATE_ENTITY,
-    TUYA_HA_TUYA_MAP,
+    TUYA_SUPPORTED_PRODUCT_CATEGORIES,
-    TUYA_HOME_MANAGER,
-    TUYA_MQTT_LISTENER,
 )
 
 _LOGGER = logging.getLogger(__name__)
 
 
+class HomeAssistantTuyaData(NamedTuple):
+    """Tuya data stored in the Home Assistant data object."""
+
+    device_listener: TuyaDeviceListener
+    device_manager: TuyaDeviceManager
+    home_manager: TuyaHomeManager
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Async setup hass config entry."""
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
+    hass.data.setdefault(DOMAIN, {})
-        TUYA_HA_TUYA_MAP: {},
-        TUYA_HA_DEVICES: set(),
-    }
 
     # Project type has been renamed to auth type in the upstream Tuya IoT SDK.
     # This migrates existing config entries to reflect that name change.
@@ -98,25 +100,28 @@ async def _init_tuya_sdk(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     tuya_mq = TuyaOpenMQ(api)
     tuya_mq.start()
 
+    device_ids: set[str] = set()
     device_manager = TuyaDeviceManager(api, tuya_mq)
-
-    # Get device list
     home_manager = TuyaHomeManager(api, tuya_mq, device_manager)
-    await hass.async_add_executor_job(home_manager.update_device_cache)
+    listener = DeviceListener(hass, device_manager, device_ids)
-    hass.data[DOMAIN][entry.entry_id][TUYA_HOME_MANAGER] = home_manager
-
-    listener = DeviceListener(hass, entry)
-    hass.data[DOMAIN][entry.entry_id][TUYA_MQTT_LISTENER] = listener
     device_manager.add_device_listener(listener)
-    hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] = device_manager
 
-    # Clean up device entities
+    hass.data[DOMAIN][entry.entry_id] = HomeAssistantTuyaData(
+        device_listener=listener,
+        device_manager=device_manager,
+        home_manager=home_manager,
+    )
+
+    # Get devices & clean up device entities
+    await hass.async_add_executor_job(home_manager.update_device_cache)
     await cleanup_device_registry(hass, device_manager)
 
-    _LOGGER.debug("init support type->%s", PLATFORMS)
+    # Register known device IDs
+    for device in device_manager.device_map.values():
+        if device.category in TUYA_SUPPORTED_PRODUCT_CATEGORIES:
+            device_ids.add(device.id)
 
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
-
     return True
 
 
@@ -134,17 +139,13 @@ async def cleanup_device_registry(
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unloading the Tuya platforms."""
-    _LOGGER.debug("integration unload")
     unload = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload:
-        device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
+        hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id]
-        device_manager.mq.stop()
+        hass_data.device_manager.mq.stop()
-        device_manager.remove_device_listener(
+        hass_data.device_manager.remove_device_listener(hass_data.device_listener)
-            hass.data[DOMAIN][entry.entry_id][TUYA_MQTT_LISTENER]
-        )
 
         hass.data[DOMAIN].pop(entry.entry_id)
-
         if not hass.data[DOMAIN]:
             hass.data.pop(DOMAIN)
 
@@ -154,14 +155,20 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 class DeviceListener(TuyaDeviceListener):
     """Device Update Listener."""
 
-    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        device_manager: TuyaDeviceManager,
+        device_ids: set[str],
+    ) -> None:
         """Init DeviceListener."""
         self.hass = hass
-        self.entry = entry
+        self.device_manager = device_manager
+        self.device_ids = device_ids
 
     def update_device(self, device: TuyaDevice) -> None:
         """Update device status."""
-        if device.id in self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_DEVICES]:
+        if device.id in self.device_ids:
             _LOGGER.debug(
                 "_update-->%s;->>%s",
                 self,
@@ -171,33 +178,14 @@ class DeviceListener(TuyaDeviceListener):
 
     def add_device(self, device: TuyaDevice) -> None:
         """Add device added listener."""
-        device_add = False
+        if device.category in TUYA_SUPPORTED_PRODUCT_CATEGORIES:
-
+            # Ensure the device isn't present stale
-        if device.category in itertools.chain(
-            *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(self.async_remove_device, device.id)
 
-            for domain, tuya_list in ha_tuya_map.items():
+            self.device_ids.add(device.id)
-                if device.category in tuya_list:
+            dispatcher_send(self.hass, TUYA_DISCOVERY_NEW, [device.id])
-                    device_add = True
-                    _LOGGER.debug(
-                        "Add device category->%s; domain-> %s",
-                        device.category,
-                        domain,
-                    )
-                    self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_DEVICES].add(
-                        device.id
-                    )
-                    dispatcher_send(
-                        self.hass, TUYA_DISCOVERY_NEW.format(domain), [device.id]
-                    )
 
-        if device_add:
+            device_manager = self.device_manager
-            device_manager = self.hass.data[DOMAIN][self.entry.entry_id][
-                TUYA_DEVICE_MANAGER
-            ]
             device_manager.mq.stop()
             tuya_mq = TuyaOpenMQ(device_manager.api)
             tuya_mq.start()
@@ -207,18 +195,16 @@ 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(self.async_remove_device, device_id)
 
     @callback
     def async_remove_device(self, device_id: str) -> None:
         """Remove device from Home Assistant."""
+        _LOGGER.debug("Remove device: %s", device_id)
         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(
+            self.device_ids.discard(device_id)
-                device_id
-            )

homeassistant/components/tuya/climate.py

@@ -8,7 +8,7 @@ from typing import Any
 
 from tuya_iot import TuyaDevice, TuyaDeviceManager
 
-from homeassistant.components.climate import DOMAIN as DEVICE_DOMAIN, ClimateEntity
+from homeassistant.components.climate import ClimateEntity
 from homeassistant.components.climate.const import (
     HVAC_MODE_AUTO,
     HVAC_MODE_COOL,
@@ -25,17 +25,11 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import TEMP_CELSIUS, TEMP_FAHRENHEIT
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
-from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import HomeAssistantTuyaData
 from .base import TuyaHaEntity
-from .const import (
+from .const import DOMAIN, TUYA_DISCOVERY_NEW
-    DOMAIN,
-    TUYA_DEVICE_MANAGER,
-    TUYA_DISCOVERY_NEW,
-    TUYA_HA_DEVICES,
-    TUYA_HA_TUYA_MAP,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -88,49 +82,24 @@ TUYA_SUPPORT_TYPE = {
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
-    """Set up tuya climate dynamically through tuya discovery."""
+    """Set up Tuya climate dynamically through Tuya discovery."""
-    _LOGGER.debug("climate init")
+    hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id]
-
-    hass.data[DOMAIN][entry.entry_id][TUYA_HA_TUYA_MAP][
-        DEVICE_DOMAIN
-    ] = TUYA_SUPPORT_TYPE
 
     @callback
-    def async_discover_device(dev_ids: list[str]) -> None:
+    def async_discover_device(device_ids: list[str]) -> None:
-        """Discover and add a discovered tuya climate."""
+        """Discover and add a discovered Tuya climate."""
-        _LOGGER.debug("climate add-> %s", dev_ids)
+        entities: list[TuyaHaClimate] = []
-        if not dev_ids:
+        for device_id in device_ids:
-            return
+            device = hass_data.device_manager.device_map[device_id]
-        entities = _setup_entities(hass, entry, dev_ids)
+            if device and device.category in TUYA_SUPPORT_TYPE:
+                entities.append(TuyaHaClimate(device, hass_data.device_manager))
         async_add_entities(entities)
 
+    async_discover_device([*hass_data.device_manager.device_map])
+
     entry.async_on_unload(
-        async_dispatcher_connect(
+        async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device)
-            hass, TUYA_DISCOVERY_NEW.format(DEVICE_DOMAIN), async_discover_device
     )
-    )
-
-    device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
-    device_ids = []
-    for (device_id, device) in device_manager.device_map.items():
-        if device.category in TUYA_SUPPORT_TYPE:
-            device_ids.append(device_id)
-    async_discover_device(device_ids)
-
-
-def _setup_entities(
-    hass: HomeAssistant, entry: ConfigEntry, device_ids: list[str]
-) -> list[Entity]:
-    """Set up Tuya Climate."""
-    device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
-    entities: list[Entity] = []
-    for device_id in device_ids:
-        device = device_manager.device_map[device_id]
-        if device is None:
-            continue
-        entities.append(TuyaHaClimate(device, device_manager))
-        hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].add(device_id)
-    return entities
 
 
 class TuyaHaClimate(TuyaHaEntity, ClimateEntity):

homeassistant/components/tuya/const.py

@@ -15,12 +15,8 @@ CONF_PASSWORD = "password"
 CONF_COUNTRY_CODE = "country_code"
 CONF_APP_TYPE = "tuya_app_type"
 
-TUYA_DISCOVERY_NEW = "tuya_discovery_new_{}"
+TUYA_DISCOVERY_NEW = "tuya_discovery_new"
-TUYA_DEVICE_MANAGER = "tuya_device_manager"
+TUYA_HA_SIGNAL_UPDATE_ENTITY = "tuya_entry_update"
-TUYA_HOME_MANAGER = "tuya_home_manager"
-TUYA_MQTT_LISTENER = "tuya_mqtt_listener"
-TUYA_HA_TUYA_MAP = "tuya_ha_tuya_map"
-TUYA_HA_DEVICES = "tuya_ha_devices"
 
 TUYA_RESPONSE_CODE = "code"
 TUYA_RESPONSE_RESULT = "result"
@@ -28,7 +24,29 @@ TUYA_RESPONSE_MSG = "msg"
 TUYA_RESPONSE_SUCCESS = "success"
 TUYA_RESPONSE_PLATFROM_URL = "platform_url"
 
-TUYA_HA_SIGNAL_UPDATE_ENTITY = "tuya_entry_update"
+TUYA_SUPPORTED_PRODUCT_CATEGORIES = (
+    "bh",  # Smart Kettle
+    "cwysj",  # Pet Water Feeder
+    "cz",  # Socket
+    "dc",  # Light string
+    "dd",  # Light strip
+    "dj",  # Light
+    "dlq",  # Breaker
+    "fs",  # Fan
+    "fs",  # Fan
+    "fwl",  # Ambient light
+    "jsq",  # Humidifier's light
+    "kg",  # Switch
+    "kj",  # Air Purifier
+    "kj",  # Air Purifier
+    "kt",  # Air conditioner
+    "pc",  # Power Strip
+    "qn",  # Heater
+    "wk",  # Thermostat
+    "xdd",  # Ceiling Light
+    "xxj",  # Diffuser
+    "xxj",  # Diffuser's light
+)
 
 TUYA_SMART_APP = "tuyaSmart"
 SMARTLIFE_APP = "smartlife"

homeassistant/components/tuya/fan.py

@@ -10,7 +10,6 @@ from tuya_iot import TuyaDevice, TuyaDeviceManager
 from homeassistant.components.fan import (
     DIRECTION_FORWARD,
     DIRECTION_REVERSE,
-    DOMAIN as DEVICE_DOMAIN,
     SUPPORT_DIRECTION,
     SUPPORT_OSCILLATE,
     SUPPORT_PRESET_MODE,
@@ -26,14 +25,9 @@ from homeassistant.util.percentage import (
     percentage_to_ordered_list_item,
 )
 
+from . import HomeAssistantTuyaData
 from .base import TuyaHaEntity
-from .const import (
+from .const import DOMAIN, TUYA_DISCOVERY_NEW
-    DOMAIN,
-    TUYA_DEVICE_MANAGER,
-    TUYA_DISCOVERY_NEW,
-    TUYA_HA_DEVICES,
-    TUYA_HA_TUYA_MAP,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -61,48 +55,23 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ):
     """Set up tuya fan dynamically through tuya discovery."""
-    _LOGGER.debug("fan init")
+    hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id]
-
-    hass.data[DOMAIN][entry.entry_id][TUYA_HA_TUYA_MAP][
-        DEVICE_DOMAIN
-    ] = TUYA_SUPPORT_TYPE
 
     @callback
-    def async_discover_device(dev_ids: list[str]) -> None:
+    def async_discover_device(device_ids: list[str]) -> None:
         """Discover and add a discovered tuya fan."""
-        _LOGGER.debug("fan add-> %s", dev_ids)
+        entities: list[TuyaHaFan] = []
-        if not dev_ids:
+        for device_id in device_ids:
-            return
+            device = hass_data.device_manager.device_map[device_id]
-        entities = _setup_entities(hass, entry, dev_ids)
+            if device and device.category in TUYA_SUPPORT_TYPE:
+                entities.append(TuyaHaFan(device, hass_data.device_manager))
         async_add_entities(entities)
 
+    async_discover_device([*hass_data.device_manager.device_map])
+
     entry.async_on_unload(
-        async_dispatcher_connect(
+        async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device)
-            hass, TUYA_DISCOVERY_NEW.format(DEVICE_DOMAIN), async_discover_device
     )
-    )
-
-    device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
-    device_ids = []
-    for (device_id, device) in device_manager.device_map.items():
-        if device.category in TUYA_SUPPORT_TYPE:
-            device_ids.append(device_id)
-    async_discover_device(device_ids)
-
-
-def _setup_entities(
-    hass: HomeAssistant, entry: ConfigEntry, device_ids: list[str]
-) -> list[TuyaHaFan]:
-    """Set up Tuya Fan."""
-    device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
-    entities = []
-    for device_id in device_ids:
-        device = device_manager.device_map[device_id]
-        if device is None:
-            continue
-        entities.append(TuyaHaFan(device, device_manager))
-        hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].add(device_id)
-    return entities
 
 
 class TuyaHaFan(TuyaHaEntity, FanEntity):

homeassistant/components/tuya/light.py

@@ -15,7 +15,6 @@ from homeassistant.components.light import (
     COLOR_MODE_COLOR_TEMP,
     COLOR_MODE_HS,
     COLOR_MODE_ONOFF,
-    DOMAIN as DEVICE_DOMAIN,
     LightEntity,
 )
 from homeassistant.config_entries import ConfigEntry
@@ -23,14 +22,9 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import HomeAssistantTuyaData
 from .base import TuyaHaEntity
-from .const import (
+from .const import DOMAIN, TUYA_DISCOVERY_NEW
-    DOMAIN,
-    TUYA_DEVICE_MANAGER,
-    TUYA_DISCOVERY_NEW,
-    TUYA_HA_DEVICES,
-    TUYA_HA_TUYA_MAP,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -84,53 +78,23 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up tuya light dynamically through tuya discovery."""
-    _LOGGER.debug("light init")
+    hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id]
-
-    hass.data[DOMAIN][entry.entry_id][TUYA_HA_TUYA_MAP][
-        DEVICE_DOMAIN
-    ] = TUYA_SUPPORT_TYPE
 
     @callback
-    def async_discover_device(dev_ids: list[str]):
+    def async_discover_device(device_ids: list[str]):
         """Discover and add a discovered tuya light."""
-        _LOGGER.debug("light add-> %s", dev_ids)
+        entities: list[TuyaHaLight] = []
-        if not dev_ids:
+        for device_id in device_ids:
-            return
+            device = hass_data.device_manager.device_map[device_id]
-        entities = _setup_entities(hass, entry, dev_ids)
+            if device and device.category in TUYA_SUPPORT_TYPE:
+                entities.append(TuyaHaLight(device, hass_data.device_manager))
         async_add_entities(entities)
 
+    async_discover_device([*hass_data.device_manager.device_map])
+
     entry.async_on_unload(
-        async_dispatcher_connect(
+        async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device)
-            hass, TUYA_DISCOVERY_NEW.format(DEVICE_DOMAIN), async_discover_device
     )
-    )
-
-    device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
-    device_ids = []
-    for (device_id, device) in device_manager.device_map.items():
-        if device.category in TUYA_SUPPORT_TYPE:
-            device_ids.append(device_id)
-    async_discover_device(device_ids)
-
-
-def _setup_entities(
-    hass, entry: ConfigEntry, device_ids: list[str]
-) -> list[TuyaHaLight]:
-    """Set up Tuya Light device."""
-    device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
-    entities = []
-    for device_id in device_ids:
-        device = device_manager.device_map[device_id]
-        if device is None:
-            continue
-
-        tuya_ha_light = TuyaHaLight(device, device_manager)
-        entities.append(tuya_ha_light)
-        hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].add(
-            tuya_ha_light.tuya_device.id
-        )
-
-    return entities
 
 
 class TuyaHaLight(TuyaHaEntity, LightEntity):

homeassistant/components/tuya/scene.py

@@ -1,7 +1,6 @@
 """Support for Tuya scenes."""
 from __future__ import annotations
 
-import logging
 from typing import Any
 
 from tuya_iot import TuyaHomeManager, TuyaScene
@@ -9,55 +8,48 @@ from tuya_iot import TuyaHomeManager, TuyaScene
 from homeassistant.components.scene import Scene
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, TUYA_HOME_MANAGER
+from . import HomeAssistantTuyaData
-
+from .const import DOMAIN
-_LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
-    """Set up tuya scenes."""
+    """Set up Tuya scenes."""
-    home_manager = hass.data[DOMAIN][entry.entry_id][TUYA_HOME_MANAGER]
+    hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id]
-    scenes = await hass.async_add_executor_job(home_manager.query_scenes)
+    scenes = await hass.async_add_executor_job(hass_data.home_manager.query_scenes)
-    async_add_entities(TuyaHAScene(home_manager, scene) for scene in scenes)
+    async_add_entities(TuyaHAScene(hass_data.home_manager, scene) for scene in scenes)
 
 
 class TuyaHAScene(Scene):
     """Tuya Scene Remote."""
 
+    _should_poll = False
+
     def __init__(self, home_manager: TuyaHomeManager, scene: TuyaScene) -> None:
         """Init Tuya Scene."""
         super().__init__()
+        self._attr_unique_id = f"tys{scene.scene_id}"
         self.home_manager = home_manager
         self.scene = scene
 
-    @property
-    def should_poll(self) -> bool:
-        """Hass should not poll."""
-        return False
-
-    @property
-    def unique_id(self) -> str | None:
-        """Return a unique ID."""
-        return f"tys{self.scene.scene_id}"
-
     @property
     def name(self) -> str | None:
         """Return Tuya scene name."""
         return self.scene.name
 
     @property
-    def device_info(self):
+    def device_info(self) -> DeviceInfo:
         """Return a device description for device registry."""
-        return {
+        return DeviceInfo(
-            "identifiers": {(DOMAIN, f"{self.unique_id}")},
+            identifiers={(DOMAIN, f"{self.unique_id}")},
-            "manufacturer": "tuya",
+            manufacturer="tuya",
-            "name": self.scene.name,
+            name=self.scene.name,
-            "model": "Tuya Scene",
+            model="Tuya Scene",
-        }
+        )
 
     @property
     def available(self) -> bool:

homeassistant/components/tuya/switch.py

@@ -6,21 +6,15 @@ from typing import Any
 
 from tuya_iot import TuyaDevice, TuyaDeviceManager
 
-from homeassistant.components.switch import DOMAIN as DEVICE_DOMAIN, SwitchEntity
+from homeassistant.components.switch import SwitchEntity
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
-from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import HomeAssistantTuyaData
 from .base import TuyaHaEntity
-from .const import (
+from .const import DOMAIN, TUYA_DISCOVERY_NEW
-    DOMAIN,
-    TUYA_DEVICE_MANAGER,
-    TUYA_DISCOVERY_NEW,
-    TUYA_HA_DEVICES,
-    TUYA_HA_TUYA_MAP,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -62,45 +56,36 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up tuya sensors dynamically through tuya discovery."""
-    _LOGGER.debug("switch init")
+    hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id]
 
-    hass.data[DOMAIN][entry.entry_id][TUYA_HA_TUYA_MAP][
+    @callback
-        DEVICE_DOMAIN
+    def async_discover_device(device_ids: list[str]) -> None:
-    ] = TUYA_SUPPORT_TYPE
-
-    async def async_discover_device(dev_ids):
         """Discover and add a discovered tuya sensor."""
-        _LOGGER.debug("switch add-> %s", dev_ids)
+        async_add_entities(
-        if not dev_ids:
+            _setup_entities(hass, entry, hass_data.device_manager, device_ids)
-            return
+        )
-        entities = _setup_entities(hass, entry, dev_ids)
+
-        async_add_entities(entities)
+    async_discover_device([*hass_data.device_manager.device_map])
 
     entry.async_on_unload(
-        async_dispatcher_connect(
+        async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device)
-            hass, TUYA_DISCOVERY_NEW.format(DEVICE_DOMAIN), async_discover_device
-        )
     )
 
-    device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
-    device_ids = []
-    for (device_id, device) in device_manager.device_map.items():
-        if device.category in TUYA_SUPPORT_TYPE:
-            device_ids.append(device_id)
-    await async_discover_device(device_ids)
 
-
+def _setup_entities(
-def _setup_entities(hass, entry: ConfigEntry, device_ids: list[str]) -> list[Entity]:
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    device_manager: TuyaDeviceManager,
+    device_ids: list[str],
+) -> list[TuyaHaSwitch]:
     """Set up Tuya Switch device."""
-    device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER]
+    entities: list[TuyaHaSwitch] = []
-    entities: list[Entity] = []
     for device_id in device_ids:
         device = device_manager.device_map[device_id]
-        if device is None:
+        if device is None or device.category not in TUYA_SUPPORT_TYPE:
             continue
 
         for function in device.function:
-            tuya_ha_switch = None
             if device.category == "kj":
                 if function in [
                     DPCODE_ANION,
@@ -110,26 +95,26 @@ def _setup_entities(hass, entry: ConfigEntry, device_ids: list[str]) -> list[Ent
                     DPCODE_UV,
                     DPCODE_WET,
                 ]:
-                    tuya_ha_switch = TuyaHaSwitch(device, device_manager, function)
+                    entities.append(TuyaHaSwitch(device, device_manager, function))
-                    # Main device switch is handled by the Fan object
+
             elif device.category == "cwysj":
-                if function in [DPCODE_FRESET, DPCODE_UV, DPCODE_PRESET, DPCODE_WRESET]:
+                if (
-                    tuya_ha_switch = TuyaHaSwitch(device, device_manager, function)
+                    function
+                    in [
+                        DPCODE_FRESET,
+                        DPCODE_UV,
+                        DPCODE_PRESET,
+                        DPCODE_WRESET,
+                    ]
+                    or function.startswith(DPCODE_SWITCH)
+                ):
+                    entities.append(TuyaHaSwitch(device, device_manager, function))
 
-                if function.startswith(DPCODE_SWITCH):
+            elif function.startswith(DPCODE_START) or function.startswith(
-                    # Main device switch
+                DPCODE_SWITCH
-                    tuya_ha_switch = TuyaHaSwitch(device, device_manager, function)
+            ):
-            else:
+                entities.append(TuyaHaSwitch(device, device_manager, function))
-                if function.startswith(DPCODE_START):
-                    tuya_ha_switch = TuyaHaSwitch(device, device_manager, function)
-                if function.startswith(DPCODE_SWITCH):
-                    tuya_ha_switch = TuyaHaSwitch(device, device_manager, function)
 
-            if tuya_ha_switch is not None:
-                entities.append(tuya_ha_switch)
-                hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].add(
-                    tuya_ha_switch.tuya_device.id
-                )
     return entities