Switch to using mqtt

custom_components/oasis_mini/button.py

@@ -1,4 +1,4 @@
-"""Oasis Mini button entity."""
+"""Oasis device button entity."""
 
 from __future__ import annotations
 
@@ -15,53 +15,54 @@ from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import OasisMiniConfigEntry
-from .entity import OasisMiniEntity
+from . import OasisDeviceConfigEntry
+from .coordinator import OasisDeviceCoordinator
+from .entity import OasisDeviceEntity
 from .helpers import add_and_play_track
-from .pyoasismini import OasisMini
-from .pyoasismini.const import TRACKS
+from .pyoasiscontrol import OasisDevice
+from .pyoasiscontrol.const import TRACKS
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: OasisMiniConfigEntry,
+    entry: OasisDeviceConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Set up Oasis Mini button using config entry."""
+    """Set up Oasis device button using config entry."""
+    coordinator: OasisDeviceCoordinator = entry.runtime_data
     async_add_entities(
-        [
-            OasisMiniButtonEntity(entry.runtime_data, descriptor)
-            for descriptor in DESCRIPTORS
-        ]
+        OasisDeviceButtonEntity(coordinator, device, descriptor)
+        for device in coordinator.data
+        for descriptor in DESCRIPTORS
     )
 
 
-async def play_random_track(device: OasisMini) -> None:
+async def play_random_track(device: OasisDevice) -> None:
     """Play random track."""
     track = random.choice(list(TRACKS))
     await add_and_play_track(device, track)
 
 
 @dataclass(frozen=True, kw_only=True)
-class OasisMiniButtonEntityDescription(ButtonEntityDescription):
-    """Oasis Mini button entity description."""
+class OasisDeviceButtonEntityDescription(ButtonEntityDescription):
+    """Oasis device button entity description."""
 
-    press_fn: Callable[[OasisMini], Awaitable[None]]
+    press_fn: Callable[[OasisDevice], Awaitable[None]]
 
 
 DESCRIPTORS = (
-    OasisMiniButtonEntityDescription(
+    OasisDeviceButtonEntityDescription(
         key="reboot",
         device_class=ButtonDeviceClass.RESTART,
         entity_category=EntityCategory.CONFIG,
         press_fn=lambda device: device.async_reboot(),
     ),
-    OasisMiniButtonEntityDescription(
+    OasisDeviceButtonEntityDescription(
         key="random_track",
         translation_key="random_track",
         press_fn=play_random_track,
     ),
-    OasisMiniButtonEntityDescription(
+    OasisDeviceButtonEntityDescription(
         key="sleep",
         translation_key="sleep",
         press_fn=lambda device: device.async_sleep(),
@@ -69,10 +70,10 @@ DESCRIPTORS = (
 )
 
 
-class OasisMiniButtonEntity(OasisMiniEntity, ButtonEntity):
-    """Oasis Mini button entity."""
+class OasisDeviceButtonEntity(OasisDeviceEntity, ButtonEntity):
+    """Oasis device button entity."""
 
-    entity_description: OasisMiniButtonEntityDescription
+    entity_description: OasisDeviceButtonEntityDescription
 
     async def async_press(self) -> None:
         """Press the button."""