Merge pull request #22 from natekspencer/dev

Fix svg

custom_components/oasis_mini/__init__.py

@@ -14,6 +14,8 @@ from .const import DOMAIN
 from .coordinator import OasisMiniCoordinator
 from .helpers import create_client
 
+type OasisMiniConfigEntry = ConfigEntry[OasisMiniCoordinator]
+
 _LOGGER = logging.getLogger(__name__)
 
 PLATFORMS = [
@@ -29,9 +31,8 @@ PLATFORMS = [
 ]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: OasisMiniConfigEntry) -> bool:
     """Set up Oasis Mini from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
     client = create_client(entry.data | entry.options)
     coordinator = OasisMiniCoordinator(hass, client)
 
@@ -62,7 +63,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         await client.session.close()
         raise ConfigEntryError("Serial number mismatch")
 
-    hass.data[DOMAIN][entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -70,15 +71,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: OasisMiniConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        await hass.data[DOMAIN][entry.entry_id].device.session.close()
-        del hass.data[DOMAIN][entry.entry_id]
-    return unload_ok
+    await entry.runtime_data.device.session.close()
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
-async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_remove_entry(hass: HomeAssistant, entry: OasisMiniConfigEntry) -> None:
     """Handle removal of an entry."""
     if entry.options:
         client = create_client(entry.data | entry.options)
@@ -86,6 +85,6 @@ async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
         await client.session.close()
 
 
-async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def update_listener(hass: HomeAssistant, entry: OasisMiniConfigEntry) -> None:
     """Handle options update."""
     await hass.config_entries.async_reload(entry.entry_id)

custom_components/oasis_mini/button.py

@@ -11,13 +11,11 @@ from homeassistant.components.button import (
     ButtonEntity,
     ButtonEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import OasisMiniCoordinator
+from . import OasisMiniConfigEntry
 from .entity import OasisMiniEntity
 from .helpers import add_and_play_track
 from .pyoasismini import OasisMini
@@ -25,13 +23,14 @@ from .pyoasismini.const import TRACKS
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: OasisMiniConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Oasis Mini button using config entry."""
-    coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
         [
-            OasisMiniButtonEntity(coordinator, entry, descriptor)
+            OasisMiniButtonEntity(entry.runtime_data, descriptor)
             for descriptor in DESCRIPTORS
         ]
     )

custom_components/oasis_mini/config_flow.py

@@ -11,7 +11,7 @@ from httpx import ConnectError, HTTPStatusError
 import voluptuous as vol
 
 from homeassistant.components import dhcp
-from homeassistant.config_entries import ConfigEntry, ConfigFlow, ConfigFlowResult
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
 from homeassistant.const import CONF_ACCESS_TOKEN, CONF_EMAIL, CONF_HOST, CONF_PASSWORD
 from homeassistant.core import callback
 from homeassistant.helpers.schema_config_entry_flow import (
@@ -21,6 +21,7 @@ from homeassistant.helpers.schema_config_entry_flow import (
     SchemaOptionsFlowHandler,
 )
 
+from . import OasisMiniConfigEntry
 from .const import DOMAIN
 from .coordinator import OasisMiniCoordinator
 from .helpers import create_client
@@ -38,9 +39,7 @@ async def cloud_login(
     handler: SchemaCommonFlowHandler, user_input: dict[str, Any]
 ) -> dict[str, Any]:
     """Cloud login."""
-    coordinator: OasisMiniCoordinator = handler.parent_handler.hass.data[DOMAIN][
-        handler.parent_handler.config_entry.entry_id
-    ]
+    coordinator: OasisMiniCoordinator = handler.parent_handler.config_entry.runtime_data
 
     try:
         await coordinator.device.async_cloud_login(
@@ -66,7 +65,9 @@ class OasisMiniConfigFlow(ConfigFlow, domain=DOMAIN):
 
     @staticmethod
     @callback
-    def async_get_options_flow(config_entry: ConfigEntry) -> SchemaOptionsFlowHandler:
+    def async_get_options_flow(
+        config_entry: OasisMiniConfigEntry,
+    ) -> SchemaOptionsFlowHandler:
         """Get the options flow for this handler."""
         return SchemaOptionsFlowHandler(config_entry, OPTIONS_FLOW)
 

custom_components/oasis_mini/entity.py

@@ -2,9 +2,6 @@
 
 from __future__ import annotations
 
-import logging
-
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, format_mac
 from homeassistant.helpers.entity import DeviceInfo, EntityDescription
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -13,8 +10,6 @@ from .const import DOMAIN
 from .coordinator import OasisMiniCoordinator
 from .pyoasismini import OasisMini
 
-_LOGGER = logging.getLogger(__name__)
-
 
 class OasisMiniEntity(CoordinatorEntity[OasisMiniCoordinator]):
     """Base class for Oasis Mini entities."""
@@ -22,10 +17,7 @@ class OasisMiniEntity(CoordinatorEntity[OasisMiniCoordinator]):
     _attr_has_entity_name = True
 
     def __init__(
-        self,
-        coordinator: OasisMiniCoordinator,
-        entry: ConfigEntry,
-        description: EntityDescription,
+        self, coordinator: OasisMiniCoordinator, description: EntityDescription
     ) -> None:
         """Construct an Oasis Mini entity."""
         super().__init__(coordinator)
@@ -37,7 +29,7 @@ class OasisMiniEntity(CoordinatorEntity[OasisMiniCoordinator]):
         self._attr_device_info = DeviceInfo(
             connections={(CONNECTION_NETWORK_MAC, format_mac(device.mac_address))},
             identifiers={(DOMAIN, serial_number)},
-            name=entry.title,
+            name=f"Oasis Mini {serial_number}",
             manufacturer="Kinetic Oasis",
             model="Oasis Mini",
             serial_number=serial_number,

custom_components/oasis_mini/image.py

@@ -3,12 +3,11 @@
 from __future__ import annotations
 
 from homeassistant.components.image import Image, ImageEntity, ImageEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import UNDEFINED
 
-from .const import DOMAIN
+from . import OasisMiniConfigEntry
 from .coordinator import OasisMiniCoordinator
 from .entity import OasisMiniEntity
 from .pyoasismini.const import TRACKS
@@ -25,13 +24,10 @@ class OasisMiniImageEntity(OasisMiniEntity, ImageEntity):
     _progress: int = 0
 
     def __init__(
-        self,
-        coordinator: OasisMiniCoordinator,
-        entry_id: str,
-        description: ImageEntityDescription,
+        self, coordinator: OasisMiniCoordinator, description: ImageEntityDescription
     ) -> None:
         """Initialize the entity."""
-        super().__init__(coordinator, entry_id, description)
+        super().__init__(coordinator, description)
         ImageEntity.__init__(self, coordinator.hass)
         self._handle_coordinator_update()
 
@@ -70,8 +66,9 @@ class OasisMiniImageEntity(OasisMiniEntity, ImageEntity):
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: OasisMiniConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Oasis Mini camera using config entry."""
-    coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([OasisMiniImageEntity(coordinator, entry, IMAGE)])
+    async_add_entities([OasisMiniImageEntity(entry.runtime_data, IMAGE)])

custom_components/oasis_mini/light.py

@@ -14,7 +14,6 @@ from homeassistant.components.light import (
     LightEntityDescription,
     LightEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.color import (
@@ -24,8 +23,7 @@ from homeassistant.util.color import (
     value_to_brightness,
 )
 
-from .const import DOMAIN
-from .coordinator import OasisMiniCoordinator
+from . import OasisMiniConfigEntry
 from .entity import OasisMiniEntity
 from .pyoasismini import LED_EFFECTS
 
@@ -70,8 +68,10 @@ class OasisMiniLightEntity(OasisMiniEntity, LightEntity):
         return self.device.brightness > 0
 
     @property
-    def rgb_color(self) -> tuple[int, int, int]:
+    def rgb_color(self) -> tuple[int, int, int] | None:
         """Return the rgb color value [int, int, int]."""
+        if not self.device.color:
+            return None
         return rgb_hex_to_rgb_list(self.device.color.replace("#", ""))
 
     @property
@@ -110,8 +110,9 @@ DESCRIPTOR = LightEntityDescription(key="led", name="LED")
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: OasisMiniConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Oasis Mini lights using config entry."""
-    coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([OasisMiniLightEntity(coordinator, entry, DESCRIPTOR)])
+    async_add_entities([OasisMiniLightEntity(entry.runtime_data, DESCRIPTOR)])

custom_components/oasis_mini/media_player.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 from datetime import datetime
-import logging
 from typing import Any
 
 from homeassistant.components.media_player import (
@@ -15,18 +14,14 @@ from homeassistant.components.media_player import (
     MediaType,
     RepeatMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ServiceValidationError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import OasisMiniCoordinator
+from . import OasisMiniConfigEntry
 from .entity import OasisMiniEntity
 from .pyoasismini.const import TRACKS
 
-_LOGGER = logging.getLogger(__name__)
-
 
 class OasisMiniMediaPlayerEntity(OasisMiniEntity, MediaPlayerEntity):
     """Oasis Mini media player entity."""
@@ -201,8 +196,9 @@ DESCRIPTOR = MediaPlayerEntityDescription(key="oasis_mini", name=None)
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: OasisMiniConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Oasis Mini media_players using config entry."""
-    coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities([OasisMiniMediaPlayerEntity(coordinator, entry, DESCRIPTOR)])
+    async_add_entities([OasisMiniMediaPlayerEntity(entry.runtime_data, DESCRIPTOR)])

custom_components/oasis_mini/number.py

@@ -7,12 +7,10 @@ from homeassistant.components.number import (
     NumberEntityDescription,
     NumberMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import OasisMiniCoordinator
+from . import OasisMiniConfigEntry
 from .entity import OasisMiniEntity
 from .pyoasismini import BALL_SPEED_MAX, BALL_SPEED_MIN, LED_SPEED_MAX, LED_SPEED_MIN
 
@@ -53,13 +51,14 @@ DESCRIPTORS = {
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: OasisMiniConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Oasis Mini numbers using config entry."""
-    coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
         [
-            OasisMiniNumberEntity(coordinator, entry, descriptor)
+            OasisMiniNumberEntity(entry.runtime_data, descriptor)
             for descriptor in DESCRIPTORS
         ]
     )

custom_components/oasis_mini/pyoasismini/__init__.py

@@ -2,10 +2,11 @@
 
 import asyncio
 import logging
-from typing import Any, Awaitable, Callable, Final
+from typing import Any, Awaitable, Final
 from urllib.parse import urljoin
 
 from aiohttp import ClientResponseError, ClientSession
+import async_timeout
 
 from .const import TRACKS
 from .utils import _bit_to_bool
@@ -32,25 +33,6 @@ AUTOPLAY_MAP = {
     "4": "30 minutes",
 }
 
-ATTRIBUTES: Final[list[tuple[str, Callable[[str], Any]]]] = [
-    ("status_code", int),  # see status code map
-    ("error", int),  # error, 0 = none, and 10 = ?, 18 = can't download?
-    ("ball_speed", int),  # 200 - 1000
-    ("playlist", lambda value: [int(track) for track in value.split(",") if track]),  # noqa: E501 # comma separated track ids
-    ("playlist_index", int),  # index of above
-    ("progress", int),  # 0 - max svg path
-    ("led_effect", str),  # led effect (code lookup)
-    ("led_color_id", str),  # led color id?
-    ("led_speed", int),  # -90 - 90
-    ("brightness", int),  # noqa: E501 # 0 - 200 in app, but seems to be 0 (off) to 304 (max), then repeats
-    ("color", str),  # hex color code
-    ("busy", _bit_to_bool),  # noqa: E501 # device is busy (downloading track, centering, software update)?
-    ("download_progress", int),  # 0 - 100%
-    ("max_brightness", int),
-    ("wifi_connected", _bit_to_bool),
-    ("repeat_playlist", _bit_to_bool),
-    ("autoplay", AUTOPLAY_MAP.get),
-]
 
 LED_EFFECTS: Final[dict[str, str]] = {
     "0": "Solid",
@@ -92,7 +74,7 @@ class OasisMini:
     autoplay: str
     brightness: int
     busy: bool
-    color: str
+    color: str | None = None
     download_progress: int
     error: int
     led_effect: str
@@ -223,14 +205,39 @@ class OasisMini:
 
     async def async_get_status(self) -> str:
         """Get the status from the device."""
-        status = await self._async_get(params={"GETSTATUS": ""})
-        _LOGGER.debug("Status: %s", status)
-        for index, value in enumerate(status.split(";")):
-            attr, func = ATTRIBUTES[index]
-            if (old_value := getattr(self, attr, None)) != (value := func(value)):
-                _LOGGER.debug("%s changed: '%s' -> '%s'", attr, old_value, value)
-                setattr(self, attr, value)
-        return status
+        raw_status = await self._async_get(params={"GETSTATUS": ""})
+        _LOGGER.debug("Status: %s", raw_status)
+        values = raw_status.split(";")
+        playlist = [int(track) for track in values[3].split(",") if track]
+        status = {
+            "status_code": int(values[0]),  # see status code map
+            "error": int(values[1]),  # noqa: E501; error, 0 = none, and 10 = ?, 18 = can't download?
+            "ball_speed": int(values[2]),  # 200 - 1000
+            "playlist": playlist,
+            "playlist_index": min(int(values[4]), len(playlist)),  # index of above
+            "progress": int(values[5]),  # 0 - max svg path
+            "led_effect": values[6],  # led effect (code lookup)
+            "led_color_id": values[7],  # led color id?
+            "led_speed": int(values[8]),  # -90 - 90
+            "brightness": int(values[9]) if values[10] else 0,  # noqa: E501; 0 - 200 in app, but seems to be 0 (off) to 304 (max), then repeats
+            "color": values[10] or None,  # hex color code
+            "busy": _bit_to_bool(values[11]),  # noqa: E501; device is busy (downloading track, centering, software update)?
+            "download_progress": int(values[12]),
+            "max_brightness": int(values[13]),
+            "wifi_connected": _bit_to_bool(values[14]),
+            "repeat_playlist": _bit_to_bool(values[15]),
+            "autoplay": AUTOPLAY_MAP.get(values[16]),
+        }
+        for key, value in status.items():
+            if (old_value := getattr(self, key, None)) != value:
+                _LOGGER.debug(
+                    "%s changed: '%s' -> '%s'",
+                    key.replace("_", " ").capitalize(),
+                    old_value,
+                    value,
+                )
+                setattr(self, key, value)
+        return raw_status
 
     async def async_move_track(self, _from: int, _to: int) -> None:
         """Move a track in the playlist."""
@@ -278,7 +285,7 @@ class OasisMini:
         if led_effect is None:
             led_effect = self.led_effect
         if color is None:
-            color = self.color
+            color = self.color or "#ffffff"
         if led_speed is None:
             led_speed = self.led_speed
         if brightness is None:

custom_components/oasis_mini/pyoasismini/utils.py

@@ -56,7 +56,7 @@ def draw_svg(track: dict, progress: int, model_id: str) -> str | None:
                         path.progress_arc_complete {{ stroke: {COLOR_DARK[1]}; }}
                         path.track {{ stroke: {COLOR_LIGHT_SHADE[1]}; }}
                         path.track_complete {{ stroke: {COLOR_MEDIUM_TINT[1]}; }}
-                    }}"""
+                    }}""".replace("\n", " ").strip()
 
                 group = SubElement(
                     svg,

custom_components/oasis_mini/select.py

@@ -6,12 +6,11 @@ from dataclasses import dataclass
 from typing import Any, Awaitable, Callable
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import OasisMiniConfigEntry
 from .coordinator import OasisMiniCoordinator
 from .entity import OasisMiniEntity
 from .pyoasismini import AUTOPLAY_MAP, OasisMini
@@ -36,11 +35,10 @@ class OasisMiniSelectEntity(OasisMiniEntity, SelectEntity):
     def __init__(
         self,
         coordinator: OasisMiniCoordinator,
-        entry: ConfigEntry[Any],
         description: EntityDescription,
     ) -> None:
         """Construct an Oasis Mini select entity."""
-        super().__init__(coordinator, entry, description)
+        super().__init__(coordinator, description)
         self._handle_coordinator_update()
 
     async def async_select_option(self, option: str) -> None:
@@ -105,13 +103,14 @@ DESCRIPTORS = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: OasisMiniConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Oasis Mini select using config entry."""
-    coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
         [
-            OasisMiniSelectEntity(coordinator, entry, descriptor)
+            OasisMiniSelectEntity(entry.runtime_data, descriptor)
             for descriptor in DESCRIPTORS
         ]
     )

custom_components/oasis_mini/sensor.py

@@ -7,29 +7,29 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import OasisMiniConfigEntry
 from .coordinator import OasisMiniCoordinator
 from .entity import OasisMiniEntity
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: OasisMiniConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Oasis Mini sensors using config entry."""
-    coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: OasisMiniCoordinator = entry.runtime_data
     entities = [
-        OasisMiniSensorEntity(coordinator, entry, descriptor)
-        for descriptor in DESCRIPTORS
+        OasisMiniSensorEntity(coordinator, descriptor) for descriptor in DESCRIPTORS
     ]
     if coordinator.device.access_token:
         entities.extend(
             [
-                OasisMiniSensorEntity(coordinator, entry, descriptor)
+                OasisMiniSensorEntity(coordinator, descriptor)
                 for descriptor in CLOUD_DESCRIPTORS
             ]
         )

custom_components/oasis_mini/switch.py

@@ -5,23 +5,22 @@
 # from typing import Any
 
 # from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-# from homeassistant.config_entries import ConfigEntry
 # from homeassistant.core import HomeAssistant
 # from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-# from .const import DOMAIN
-# from .coordinator import OasisMiniCoordinator
+# from . import OasisMiniConfigEntry
 # from .entity import OasisMiniEntity
 
 
 # async def async_setup_entry(
-#     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+#     hass: HomeAssistant,
+#     entry: OasisMiniConfigEntry,
+#     async_add_entities: AddEntitiesCallback,
 # ) -> None:
 #     """Set up Oasis Mini switchs using config entry."""
-#     coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
 #     async_add_entities(
 #         [
-#             OasisMiniSwitchEntity(coordinator, entry, descriptor)
+#             OasisMiniSwitchEntity(entry.runtime_data, descriptor)
 #             for descriptor in DESCRIPTORS
 #         ]
 #     )

custom_components/oasis_mini/update.py

@@ -12,11 +12,10 @@ from homeassistant.components.update import (
     UpdateEntityDescription,
     UpdateEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import OasisMiniConfigEntry
 from .coordinator import OasisMiniCoordinator
 from .entity import OasisMiniEntity
 
@@ -26,14 +25,14 @@ SCAN_INTERVAL = timedelta(hours=6)
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: OasisMiniConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Oasis Mini updates using config entry."""
-    coordinator: OasisMiniCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: OasisMiniCoordinator = entry.runtime_data
     if coordinator.device.access_token:
-        async_add_entities(
-            [OasisMiniUpdateEntity(coordinator, entry, DESCRIPTOR)], True
-        )
+        async_add_entities([OasisMiniUpdateEntity(coordinator, DESCRIPTOR)], True)
 
 
 DESCRIPTOR = UpdateEntityDescription(