From 63cb79ec29899a6efdf84d8c076b8d232f490f94 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 14 Feb 2022 09:49:43 -0600
Subject: [PATCH] Add support for setting the effect speed in WiZ (#66457)

---
 homeassistant/components/wiz/__init__.py |  2 +-
 homeassistant/components/wiz/number.py   | 58 ++++++++++++++++++++++++
 tests/components/wiz/__init__.py         |  3 +-
 tests/components/wiz/test_number.py      | 32 +++++++++++++
 4 files changed, 93 insertions(+), 2 deletions(-)
 create mode 100644 homeassistant/components/wiz/number.py
 create mode 100644 tests/components/wiz/test_number.py

diff --git a/homeassistant/components/wiz/__init__.py b/homeassistant/components/wiz/__init__.py
index 40dc4cf70d1..1bed875d02f 100644
--- a/homeassistant/components/wiz/__init__.py
+++ b/homeassistant/components/wiz/__init__.py
@@ -30,7 +30,7 @@ from .models import WizData
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS = [Platform.BINARY_SENSOR, Platform.LIGHT, Platform.SWITCH]
+PLATFORMS = [Platform.BINARY_SENSOR, Platform.LIGHT, Platform.NUMBER, Platform.SWITCH]
 
 REQUEST_REFRESH_DELAY = 0.35
 
diff --git a/homeassistant/components/wiz/number.py b/homeassistant/components/wiz/number.py
new file mode 100644
index 00000000000..eed9bd92803
--- /dev/null
+++ b/homeassistant/components/wiz/number.py
@@ -0,0 +1,58 @@
+"""Support for WiZ effect speed numbers."""
+from __future__ import annotations
+
+from pywizlight.bulblibrary import BulbClass
+
+from homeassistant.components.number import NumberEntity, NumberMode
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .const import DOMAIN
+from .entity import WizEntity
+from .models import WizData
+
+EFFECT_SPEED_UNIQUE_ID = "{}_effect_speed"
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up the wiz speed number."""
+    wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
+    if wiz_data.bulb.bulbtype.bulb_type != BulbClass.SOCKET:
+        async_add_entities([WizSpeedNumber(wiz_data, entry.title)])
+
+
+class WizSpeedNumber(WizEntity, NumberEntity):
+    """Defines a WiZ speed number."""
+
+    _attr_min_value = 10
+    _attr_max_value = 200
+    _attr_step = 1
+    _attr_mode = NumberMode.SLIDER
+    _attr_icon = "mdi:speedometer"
+
+    def __init__(self, wiz_data: WizData, name: str) -> None:
+        """Initialize an WiZ device."""
+        super().__init__(wiz_data, name)
+        self._attr_unique_id = EFFECT_SPEED_UNIQUE_ID.format(self._device.mac)
+        self._attr_name = f"{name} Effect Speed"
+        self._async_update_attrs()
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        return super().available and self._device.state.get_speed() is not None
+
+    @callback
+    def _async_update_attrs(self) -> None:
+        """Handle updating _attr values."""
+        self._attr_value = self._device.state.get_speed()
+
+    async def async_set_value(self, value: float) -> None:
+        """Set the speed value."""
+        await self._device.set_speed(int(value))
+        await self.coordinator.async_request_refresh()
diff --git a/tests/components/wiz/__init__.py b/tests/components/wiz/__init__.py
index e553593bf2f..8d187e9b476 100644
--- a/tests/components/wiz/__init__.py
+++ b/tests/components/wiz/__init__.py
@@ -170,6 +170,7 @@ def _mocked_wizlight(device, extended_white_range, bulb_type) -> wizlight:
     bulb.getSupportedScenes = AsyncMock(return_value=list(SCENES))
     bulb.start_push = AsyncMock(side_effect=_save_setup_callback)
     bulb.async_close = AsyncMock()
+    bulb.set_speed = AsyncMock()
     bulb.state = FAKE_STATE
     bulb.mac = FAKE_MAC
     bulb.bulbtype = bulb_type or FAKE_DIMMABLE_BULB
@@ -223,6 +224,6 @@ async def async_setup_integration(
 async def async_push_update(hass, device, params):
     """Push an update to the device."""
     device.state = PilotParser(params)
-    device.status = params["state"]
+    device.status = params.get("state")
     device.push_callback(device.state)
     await hass.async_block_till_done()
diff --git a/tests/components/wiz/test_number.py b/tests/components/wiz/test_number.py
new file mode 100644
index 00000000000..a1ab5e6bbae
--- /dev/null
+++ b/tests/components/wiz/test_number.py
@@ -0,0 +1,32 @@
+"""Tests for the number platform."""
+
+from homeassistant.components.number import DOMAIN as NUMBER_DOMAIN
+from homeassistant.components.number.const import ATTR_VALUE, SERVICE_SET_VALUE
+from homeassistant.const import ATTR_ENTITY_ID, STATE_UNAVAILABLE
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+
+from . import FAKE_MAC, async_push_update, async_setup_integration
+
+
+async def test_speed_operation(hass: HomeAssistant) -> None:
+    """Test changing a speed."""
+    bulb, _ = await async_setup_integration(hass)
+    await async_push_update(hass, bulb, {"mac": FAKE_MAC})
+    entity_id = "number.mock_title_effect_speed"
+    entity_registry = er.async_get(hass)
+    assert entity_registry.async_get(entity_id).unique_id == f"{FAKE_MAC}_effect_speed"
+    assert hass.states.get(entity_id).state == STATE_UNAVAILABLE
+
+    await async_push_update(hass, bulb, {"mac": FAKE_MAC, "speed": 50})
+    assert hass.states.get(entity_id).state == "50"
+
+    await hass.services.async_call(
+        NUMBER_DOMAIN,
+        SERVICE_SET_VALUE,
+        {ATTR_ENTITY_ID: entity_id, ATTR_VALUE: 30},
+        blocking=True,
+    )
+    bulb.set_speed.assert_called_with(30)
+    await async_push_update(hass, bulb, {"mac": FAKE_MAC, "speed": 30})
+    assert hass.states.get(entity_id).state == "30"