Add Switch entity to SleepIQ (#66966)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/sleepiq/__init__.py

@@ -18,11 +18,15 @@ import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.typing import ConfigType
 
 from .const import DOMAIN
-from .coordinator import SleepIQDataUpdateCoordinator
+from .coordinator import (
+    SleepIQData,
+    SleepIQDataUpdateCoordinator,
+    SleepIQPauseUpdateCoordinator,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS = [Platform.BINARY_SENSOR, Platform.BUTTON, Platform.SENSOR]
+PLATFORMS = [Platform.BINARY_SENSOR, Platform.BUTTON, Platform.SENSOR, Platform.SWITCH]
 
 CONFIG_SCHEMA = vol.Schema(
     {
@@ -77,11 +81,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady(str(err) or "Error reading from SleepIQ API") from err
 
     coordinator = SleepIQDataUpdateCoordinator(hass, gateway, email)
+    pause_coordinator = SleepIQPauseUpdateCoordinator(hass, gateway, email)
 
     # Call the SleepIQ API to refresh data
     await coordinator.async_config_entry_first_refresh()
+    await pause_coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = SleepIQData(
+        data_coordinator=coordinator,
+        pause_coordinator=pause_coordinator,
+        client=gateway,
+    )
 
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 

homeassistant/components/sleepiq/binary_sensor.py

@@ -11,7 +11,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import DOMAIN, ICON_EMPTY, ICON_OCCUPIED, IS_IN_BED
-from .coordinator import SleepIQDataUpdateCoordinator
+from .coordinator import SleepIQData
 from .entity import SleepIQSensor
 
 
@@ -21,10 +21,10 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the SleepIQ bed binary sensors."""
-    coordinator: SleepIQDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    data: SleepIQData = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
-        IsInBedBinarySensor(coordinator, bed, sleeper)
-        for bed in coordinator.client.beds.values()
+        IsInBedBinarySensor(data.data_coordinator, bed, sleeper)
+        for bed in data.client.beds.values()
         for sleeper in bed.sleepers
     )
 

homeassistant/components/sleepiq/button.py

@@ -13,7 +13,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .coordinator import SleepIQDataUpdateCoordinator
+from .coordinator import SleepIQData
 from .entity import SleepIQEntity
 
 
@@ -53,11 +53,11 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the sleep number buttons."""
-    coordinator: SleepIQDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    data: SleepIQData = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
         SleepNumberButton(bed, ed)
-        for bed in coordinator.client.beds.values()
+        for bed in data.client.beds.values()
         for ed in ENTITY_DESCRIPTIONS
     )
 

homeassistant/components/sleepiq/coordinator.py

@@ -1,4 +1,6 @@
 """Coordinator for SleepIQ."""
+import asyncio
+from dataclasses import dataclass
 from datetime import timedelta
 import logging
 
@@ -10,9 +12,10 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 _LOGGER = logging.getLogger(__name__)
 
 UPDATE_INTERVAL = timedelta(seconds=60)
+LONGER_UPDATE_INTERVAL = timedelta(minutes=5)
 
 
-class SleepIQDataUpdateCoordinator(DataUpdateCoordinator[dict[str, dict]]):
+class SleepIQDataUpdateCoordinator(DataUpdateCoordinator[None]):
     """SleepIQ data update coordinator."""
 
     def __init__(
@@ -26,7 +29,42 @@ class SleepIQDataUpdateCoordinator(DataUpdateCoordinator[dict[str, dict]]):
             hass,
             _LOGGER,
             name=f"{username}@SleepIQ",
-            update_method=client.fetch_bed_statuses,
             update_interval=UPDATE_INTERVAL,
         )
         self.client = client
+
+    async def _async_update_data(self) -> None:
+        await self.client.fetch_bed_statuses()
+
+
+class SleepIQPauseUpdateCoordinator(DataUpdateCoordinator[None]):
+    """SleepIQ data update coordinator."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        client: AsyncSleepIQ,
+        username: str,
+    ) -> None:
+        """Initialize coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=f"{username}@SleepIQPause",
+            update_interval=LONGER_UPDATE_INTERVAL,
+        )
+        self.client = client
+
+    async def _async_update_data(self) -> None:
+        await asyncio.gather(
+            *[bed.fetch_pause_mode() for bed in self.client.beds.values()]
+        )
+
+
+@dataclass
+class SleepIQData:
+    """Data for the sleepiq integration."""
+
+    data_coordinator: SleepIQDataUpdateCoordinator
+    pause_coordinator: SleepIQPauseUpdateCoordinator
+    client: AsyncSleepIQ

homeassistant/components/sleepiq/entity.py

@@ -14,18 +14,23 @@ from homeassistant.helpers.update_coordinator import (
 from .const import ICON_OCCUPIED, SENSOR_TYPES
 
 
+def device_from_bed(bed: SleepIQBed) -> DeviceInfo:
+    """Create a device given a bed."""
+    return DeviceInfo(
+        connections={(device_registry.CONNECTION_NETWORK_MAC, bed.mac_addr)},
+        manufacturer="SleepNumber",
+        name=bed.name,
+        model=bed.model,
+    )
+
+
 class SleepIQEntity(Entity):
     """Implementation of a SleepIQ entity."""
 
     def __init__(self, bed: SleepIQBed) -> None:
         """Initialize the SleepIQ entity."""
         self.bed = bed
-        self._attr_device_info = DeviceInfo(
-            connections={(device_registry.CONNECTION_NETWORK_MAC, bed.mac_addr)},
-            manufacturer="SleepNumber",
-            name=bed.name,
-            model=bed.model,
-        )
+        self._attr_device_info = device_from_bed(bed)
 
 
 class SleepIQSensor(CoordinatorEntity):
@@ -44,12 +49,7 @@ class SleepIQSensor(CoordinatorEntity):
         super().__init__(coordinator)
         self.sleeper = sleeper
         self.bed = bed
-        self._attr_device_info = DeviceInfo(
-            connections={(device_registry.CONNECTION_NETWORK_MAC, bed.mac_addr)},
-            manufacturer="SleepNumber",
-            name=bed.name,
-            model=bed.model,
-        )
+        self._attr_device_info = device_from_bed(bed)
 
         self._attr_name = f"SleepNumber {bed.name} {sleeper.name} {SENSOR_TYPES[name]}"
         self._attr_unique_id = f"{bed.id}_{sleeper.name}_{name}"
@@ -65,3 +65,19 @@ class SleepIQSensor(CoordinatorEntity):
     @abstractmethod
     def _async_update_attrs(self) -> None:
         """Update sensor attributes."""
+
+
+class SleepIQBedCoordinator(CoordinatorEntity):
+    """Implementation of a SleepIQ sensor."""
+
+    _attr_icon = ICON_OCCUPIED
+
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator,
+        bed: SleepIQBed,
+    ) -> None:
+        """Initialize the SleepIQ sensor entity."""
+        super().__init__(coordinator)
+        self.bed = bed
+        self._attr_device_info = device_from_bed(bed)

homeassistant/components/sleepiq/sensor.py

@@ -10,7 +10,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import DOMAIN, SLEEP_NUMBER
-from .coordinator import SleepIQDataUpdateCoordinator
+from .coordinator import SleepIQData
 from .entity import SleepIQSensor
 
 
@@ -20,10 +20,10 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the SleepIQ bed sensors."""
-    coordinator: SleepIQDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    data: SleepIQData = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
-        SleepNumberSensorEntity(coordinator, bed, sleeper)
-        for bed in coordinator.client.beds.values()
+        SleepNumberSensorEntity(data.data_coordinator, bed, sleeper)
+        for bed in data.client.beds.values()
         for sleeper in bed.sleepers
     )
 

homeassistant/components/sleepiq/switch.py

@@ -0,0 +1,53 @@
+"""Support for SleepIQ switches."""
+from __future__ import annotations
+
+from typing import Any
+
+from asyncsleepiq import SleepIQBed
+
+from homeassistant.components.switch import SwitchEntity
+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 SleepIQData, SleepIQPauseUpdateCoordinator
+from .entity import SleepIQBedCoordinator
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up the sleep number switches."""
+    data: SleepIQData = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities(
+        SleepNumberPrivateSwitch(data.pause_coordinator, bed)
+        for bed in data.client.beds.values()
+    )
+
+
+class SleepNumberPrivateSwitch(SleepIQBedCoordinator, SwitchEntity):
+    """Representation of SleepIQ privacy mode."""
+
+    def __init__(
+        self, coordinator: SleepIQPauseUpdateCoordinator, bed: SleepIQBed
+    ) -> None:
+        """Initialize the switch."""
+        super().__init__(coordinator, bed)
+        self._attr_name = f"SleepNumber {bed.name} Pause Mode"
+        self._attr_unique_id = f"{bed.id}-pause-mode"
+
+    @property
+    def is_on(self) -> bool:
+        """Return whether the switch is on or off."""
+        return bool(self.bed.paused)
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn on switch."""
+        await self.bed.set_pause_mode(True)
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn off switch."""
+        await self.bed.set_pause_mode(False)

tests/components/sleepiq/test_switch.py

@@ -0,0 +1,69 @@
+"""The tests for SleepIQ switch platform."""
+from homeassistant.components.sleepiq.coordinator import LONGER_UPDATE_INTERVAL
+from homeassistant.components.switch import DOMAIN
+from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON
+from homeassistant.helpers import entity_registry as er
+from homeassistant.util.dt import utcnow
+
+from tests.common import async_fire_time_changed
+from tests.components.sleepiq.conftest import (
+    BED_ID,
+    BED_NAME,
+    BED_NAME_LOWER,
+    setup_platform,
+)
+
+
+async def test_setup(hass, mock_asyncsleepiq):
+    """Test for successfully setting up the SleepIQ platform."""
+    entry = await setup_platform(hass, DOMAIN)
+    entity_registry = er.async_get(hass)
+
+    assert len(entity_registry.entities) == 1
+
+    entry = entity_registry.async_get(f"switch.sleepnumber_{BED_NAME_LOWER}_pause_mode")
+    assert entry
+    assert entry.original_name == f"SleepNumber {BED_NAME} Pause Mode"
+    assert entry.unique_id == f"{BED_ID}-pause-mode"
+
+
+async def test_switch_set_states(hass, mock_asyncsleepiq):
+    """Test button press."""
+    await setup_platform(hass, DOMAIN)
+
+    await hass.services.async_call(
+        DOMAIN,
+        "turn_off",
+        {ATTR_ENTITY_ID: f"switch.sleepnumber_{BED_NAME_LOWER}_pause_mode"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    mock_asyncsleepiq.beds[BED_ID].set_pause_mode.assert_called_with(False)
+
+    await hass.services.async_call(
+        DOMAIN,
+        "turn_on",
+        {ATTR_ENTITY_ID: f"switch.sleepnumber_{BED_NAME_LOWER}_pause_mode"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    mock_asyncsleepiq.beds[BED_ID].set_pause_mode.assert_called_with(True)
+
+
+async def test_switch_get_states(hass, mock_asyncsleepiq):
+    """Test button press."""
+    await setup_platform(hass, DOMAIN)
+
+    assert (
+        hass.states.get(f"switch.sleepnumber_{BED_NAME_LOWER}_pause_mode").state
+        == STATE_OFF
+    )
+    mock_asyncsleepiq.beds[BED_ID].paused = True
+
+    async_fire_time_changed(hass, utcnow() + LONGER_UPDATE_INTERVAL)
+    await hass.async_block_till_done()
+
+    assert (
+        hass.states.get(f"switch.sleepnumber_{BED_NAME_LOWER}_pause_mode").state
+        == STATE_ON
+    )