Restructure WLED integration (#51667)

homeassistant/components/wled/__init__.py

@@ -1,44 +1,21 @@
 """Support for WLED."""
 from __future__ import annotations
 
-from datetime import timedelta
-import logging
-
-from wled import WLED, Device as WLEDDevice, WLEDConnectionError, WLEDError
-
 from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
-    ATTR_IDENTIFIERS,
-    ATTR_MANUFACTURER,
-    ATTR_MODEL,
-    ATTR_NAME,
-    ATTR_SW_VERSION,
-    CONF_HOST,
-)
+from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.entity import DeviceInfo
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-    UpdateFailed,
-)
 
 from .const import DOMAIN
+from .coordinator import WLEDDataUpdateCoordinator
 
-SCAN_INTERVAL = timedelta(seconds=5)
 PLATFORMS = (LIGHT_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN)
 
-_LOGGER = logging.getLogger(__name__)
-
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up WLED from a config entry."""
-
-    # Create WLED instance for this entry
     coordinator = WLEDDataUpdateCoordinator(hass, host=entry.data[CONF_HOST])
     await coordinator.async_config_entry_first_refresh()
 
@@ -59,85 +36,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload WLED config entry."""
-
-    # Unload entities for this entry/device.
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
         del hass.data[DOMAIN][entry.entry_id]
-
-    if not hass.data[DOMAIN]:
-        del hass.data[DOMAIN]
-
     return unload_ok
-
-
-def wled_exception_handler(func):
-    """Decorate WLED calls to handle WLED exceptions.
-
-    A decorator that wraps the passed in function, catches WLED errors,
-    and handles the availability of the device in the data coordinator.
-    """
-
-    async def handler(self, *args, **kwargs):
-        try:
-            await func(self, *args, **kwargs)
-            self.coordinator.update_listeners()
-
-        except WLEDConnectionError as error:
-            _LOGGER.error("Error communicating with API: %s", error)
-            self.coordinator.last_update_success = False
-            self.coordinator.update_listeners()
-
-        except WLEDError as error:
-            _LOGGER.error("Invalid response from API: %s", error)
-
-    return handler
-
-
-class WLEDDataUpdateCoordinator(DataUpdateCoordinator[WLEDDevice]):
-    """Class to manage fetching WLED data from single endpoint."""
-
-    def __init__(
-        self,
-        hass: HomeAssistant,
-        *,
-        host: str,
-    ) -> None:
-        """Initialize global WLED data updater."""
-        self.wled = WLED(host, session=async_get_clientsession(hass))
-
-        super().__init__(
-            hass,
-            _LOGGER,
-            name=DOMAIN,
-            update_interval=SCAN_INTERVAL,
-        )
-
-    def update_listeners(self) -> None:
-        """Call update on all listeners."""
-        for update_callback in self._listeners:
-            update_callback()
-
-    async def _async_update_data(self) -> WLEDDevice:
-        """Fetch data from WLED."""
-        try:
-            return await self.wled.update(full_update=not self.last_update_success)
-        except WLEDError as error:
-            raise UpdateFailed(f"Invalid response from API: {error}") from error
-
-
-class WLEDEntity(CoordinatorEntity):
-    """Defines a base WLED entity."""
-
-    coordinator: WLEDDataUpdateCoordinator
-
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return device information about this WLED device."""
-        return {
-            ATTR_IDENTIFIERS: {(DOMAIN, self.coordinator.data.info.mac_address)},
-            ATTR_NAME: self.coordinator.data.info.name,
-            ATTR_MANUFACTURER: self.coordinator.data.info.brand,
-            ATTR_MODEL: self.coordinator.data.info.product,
-            ATTR_SW_VERSION: self.coordinator.data.info.version,
-        }

homeassistant/components/wled/const.py

@@ -1,8 +1,13 @@
 """Constants for the WLED integration."""
+from datetime import timedelta
+import logging
 
 # Integration domain
 DOMAIN = "wled"
 
+LOGGER = logging.getLogger(__package__)
+SCAN_INTERVAL = timedelta(seconds=5)
+
 # Attributes
 ATTR_COLOR_PRIMARY = "color_primary"
 ATTR_DURATION = "duration"

homeassistant/components/wled/coordinator.py

@@ -0,0 +1,41 @@
+"""DataUpdateCoordinator for WLED."""
+
+from wled import WLED, Device as WLEDDevice, WLEDError
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN, LOGGER, SCAN_INTERVAL
+
+
+class WLEDDataUpdateCoordinator(DataUpdateCoordinator[WLEDDevice]):
+    """Class to manage fetching WLED data from single endpoint."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        *,
+        host: str,
+    ) -> None:
+        """Initialize global WLED data updater."""
+        self.wled = WLED(host, session=async_get_clientsession(hass))
+
+        super().__init__(
+            hass,
+            LOGGER,
+            name=DOMAIN,
+            update_interval=SCAN_INTERVAL,
+        )
+
+    def update_listeners(self) -> None:
+        """Call update on all listeners."""
+        for update_callback in self._listeners:
+            update_callback()
+
+    async def _async_update_data(self) -> WLEDDevice:
+        """Fetch data from WLED."""
+        try:
+            return await self.wled.update(full_update=not self.last_update_success)
+        except WLEDError as error:
+            raise UpdateFailed(f"Invalid response from API: {error}") from error

homeassistant/components/wled/helpers.py

@@ -0,0 +1,28 @@
+"""Helpers for WLED."""
+
+from wled import WLEDConnectionError, WLEDError
+
+from .const import LOGGER
+
+
+def wled_exception_handler(func):
+    """Decorate WLED calls to handle WLED exceptions.
+
+    A decorator that wraps the passed in function, catches WLED errors,
+    and handles the availability of the device in the data coordinator.
+    """
+
+    async def handler(self, *args, **kwargs):
+        try:
+            await func(self, *args, **kwargs)
+            self.coordinator.update_listeners()
+
+        except WLEDConnectionError as error:
+            LOGGER.error("Error communicating with API: %s", error)
+            self.coordinator.last_update_success = False
+            self.coordinator.update_listeners()
+
+        except WLEDError as error:
+            LOGGER.error("Invalid response from API: %s", error)
+
+    return handler

homeassistant/components/wled/light.py

@@ -27,7 +27,6 @@ from homeassistant.helpers.entity_registry import (
     async_get_registry as async_get_entity_registry,
 )
 
-from . import WLEDDataUpdateCoordinator, WLEDEntity, wled_exception_handler
 from .const import (
     ATTR_COLOR_PRIMARY,
     ATTR_INTENSITY,
@@ -42,6 +41,9 @@ from .const import (
     SERVICE_EFFECT,
     SERVICE_PRESET,
 )
+from .coordinator import WLEDDataUpdateCoordinator
+from .helpers import wled_exception_handler
+from .models import WLEDEntity
 
 PARALLEL_UPDATES = 1
 

homeassistant/components/wled/models.py

@@ -0,0 +1,30 @@
+"""Models for WLED."""
+from homeassistant.const import (
+    ATTR_IDENTIFIERS,
+    ATTR_MANUFACTURER,
+    ATTR_MODEL,
+    ATTR_NAME,
+    ATTR_SW_VERSION,
+)
+from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .const import DOMAIN
+from .coordinator import WLEDDataUpdateCoordinator
+
+
+class WLEDEntity(CoordinatorEntity):
+    """Defines a base WLED entity."""
+
+    coordinator: WLEDDataUpdateCoordinator
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information about this WLED device."""
+        return {
+            ATTR_IDENTIFIERS: {(DOMAIN, self.coordinator.data.info.mac_address)},
+            ATTR_NAME: self.coordinator.data.info.name,
+            ATTR_MANUFACTURER: self.coordinator.data.info.brand,
+            ATTR_MODEL: self.coordinator.data.info.product,
+            ATTR_SW_VERSION: self.coordinator.data.info.version,
+        }

homeassistant/components/wled/sensor.py

@@ -17,8 +17,9 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.dt import utcnow
 
-from . import WLEDDataUpdateCoordinator, WLEDEntity
 from .const import ATTR_LED_COUNT, ATTR_MAX_POWER, CURRENT_MA, DOMAIN
+from .coordinator import WLEDDataUpdateCoordinator
+from .models import WLEDEntity
 
 
 async def async_setup_entry(

homeassistant/components/wled/switch.py

@@ -8,7 +8,6 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import WLEDDataUpdateCoordinator, WLEDEntity, wled_exception_handler
 from .const import (
     ATTR_DURATION,
     ATTR_FADE,
@@ -16,6 +15,9 @@ from .const import (
     ATTR_UDP_PORT,
     DOMAIN,
 )
+from .coordinator import WLEDDataUpdateCoordinator
+from .helpers import wled_exception_handler
+from .models import WLEDEntity
 
 PARALLEL_UPDATES = 1
 

tests/components/wled/test_config_flow.py

@@ -91,7 +91,10 @@ async def test_full_zeroconf_flow_implementation(
     assert result2["data"][CONF_MAC] == "aabbccddeeff"
 
 
-@patch("homeassistant.components.wled.WLED.update", side_effect=WLEDConnectionError)
+@patch(
+    "homeassistant.components.wled.coordinator.WLED.update",
+    side_effect=WLEDConnectionError,
+)
 async def test_connection_error(
     update_mock: MagicMock, hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
 ) -> None:
@@ -109,7 +112,10 @@ async def test_connection_error(
     assert result.get("errors") == {"base": "cannot_connect"}
 
 
-@patch("homeassistant.components.wled.WLED.update", side_effect=WLEDConnectionError)
+@patch(
+    "homeassistant.components.wled.coordinator.WLED.update",
+    side_effect=WLEDConnectionError,
+)
 async def test_zeroconf_connection_error(
     update_mock: MagicMock, hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
 ) -> None:
@@ -126,7 +132,10 @@ async def test_zeroconf_connection_error(
     assert result.get("reason") == "cannot_connect"
 
 
-@patch("homeassistant.components.wled.WLED.update", side_effect=WLEDConnectionError)
+@patch(
+    "homeassistant.components.wled.coordinator.WLED.update",
+    side_effect=WLEDConnectionError,
+)
 async def test_zeroconf_confirm_connection_error(
     update_mock: MagicMock, hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
 ) -> None:

tests/components/wled/test_init.py

@@ -11,7 +11,10 @@ from tests.components.wled import init_integration
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 
-@patch("homeassistant.components.wled.WLED.update", side_effect=WLEDConnectionError)
+@patch(
+    "homeassistant.components.wled.coordinator.WLED.update",
+    side_effect=WLEDConnectionError,
+)
 async def test_config_entry_not_ready(
     mock_update: MagicMock, hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
 ) -> None:

tests/components/wled/test_light.py

@@ -13,7 +13,6 @@ from homeassistant.components.light import (
     ATTR_TRANSITION,
     DOMAIN as LIGHT_DOMAIN,
 )
-from homeassistant.components.wled import SCAN_INTERVAL
 from homeassistant.components.wled.const import (
     ATTR_INTENSITY,
     ATTR_PALETTE,
@@ -22,6 +21,7 @@ from homeassistant.components.wled.const import (
     ATTR_REVERSE,
     ATTR_SPEED,
     DOMAIN,
+    SCAN_INTERVAL,
     SERVICE_EFFECT,
     SERVICE_PRESET,
 )
@@ -228,7 +228,7 @@ async def test_dynamically_handle_segments(
 
     # Test removal if segment went missing, including the master entity
     with patch(
-        "homeassistant.components.wled.WLED.update",
+        "homeassistant.components.wled.coordinator.WLED.update",
         return_value=device,
     ):
         async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL)
@@ -257,7 +257,7 @@ async def test_single_segment_behavior(
 
     # Test absent master
     with patch(
-        "homeassistant.components.wled.WLED.update",
+        "homeassistant.components.wled.coordinator.WLED.update",
         return_value=device,
     ):
         async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL)
@@ -273,7 +273,7 @@ async def test_single_segment_behavior(
     device.state.brightness = 100
     device.state.segments[0].brightness = 255
     with patch(
-        "homeassistant.components.wled.WLED.update",
+        "homeassistant.components.wled.coordinator.WLED.update",
         return_value=device,
     ):
         async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL)
@@ -286,7 +286,7 @@ async def test_single_segment_behavior(
     # Test segment is off when master is off
     device.state.on = False
     with patch(
-        "homeassistant.components.wled.WLED.update",
+        "homeassistant.components.wled.coordinator.WLED.update",
         return_value=device,
     ):
         async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL)
@@ -336,7 +336,7 @@ async def test_light_error(
     aioclient_mock.post("http://192.168.1.123:80/json/state", text="", status=400)
     await init_integration(hass, aioclient_mock)
 
-    with patch("homeassistant.components.wled.WLED.update"):
+    with patch("homeassistant.components.wled.coordinator.WLED.update"):
         await hass.services.async_call(
             LIGHT_DOMAIN,
             SERVICE_TURN_OFF,
@@ -356,8 +356,9 @@ async def test_light_connection_error(
     """Test error handling of the WLED switches."""
     await init_integration(hass, aioclient_mock)
 
-    with patch("homeassistant.components.wled.WLED.update"), patch(
-        "homeassistant.components.wled.WLED.segment", side_effect=WLEDConnectionError
+    with patch("homeassistant.components.wled.coordinator.WLED.update"), patch(
+        "homeassistant.components.wled.coordinator.WLED.segment",
+        side_effect=WLEDConnectionError,
     ):
         await hass.services.async_call(
             LIGHT_DOMAIN,
@@ -532,7 +533,7 @@ async def test_effect_service_error(
     aioclient_mock.post("http://192.168.1.123:80/json/state", text="", status=400)
     await init_integration(hass, aioclient_mock)
 
-    with patch("homeassistant.components.wled.WLED.update"):
+    with patch("homeassistant.components.wled.coordinator.WLED.update"):
         await hass.services.async_call(
             DOMAIN,
             SERVICE_EFFECT,
@@ -575,7 +576,7 @@ async def test_preset_service_error(
     aioclient_mock.post("http://192.168.1.123:80/json/state", text="", status=400)
     await init_integration(hass, aioclient_mock)
 
-    with patch("homeassistant.components.wled.WLED.update"):
+    with patch("homeassistant.components.wled.coordinator.WLED.update"):
         await hass.services.async_call(
             DOMAIN,
             SERVICE_PRESET,

tests/components/wled/test_switch.py

@@ -144,7 +144,7 @@ async def test_switch_error(
     aioclient_mock.post("http://192.168.1.123:80/json/state", text="", status=400)
     await init_integration(hass, aioclient_mock)
 
-    with patch("homeassistant.components.wled.WLED.update"):
+    with patch("homeassistant.components.wled.coordinator.WLED.update"):
         await hass.services.async_call(
             SWITCH_DOMAIN,
             SERVICE_TURN_ON,
@@ -164,8 +164,9 @@ async def test_switch_connection_error(
     """Test error handling of the WLED switches."""
     await init_integration(hass, aioclient_mock)
 
-    with patch("homeassistant.components.wled.WLED.update"), patch(
-        "homeassistant.components.wled.WLED.nightlight", side_effect=WLEDConnectionError
+    with patch("homeassistant.components.wled.coordinator.WLED.update"), patch(
+        "homeassistant.components.wled.coordinator.WLED.nightlight",
+        side_effect=WLEDConnectionError,
     ):
         await hass.services.async_call(
             SWITCH_DOMAIN,