From fc4cb743bdb046c184df6959d54ffecce949009f Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 28 Feb 2022 18:50:26 -1000
Subject: [PATCH] Add Signal Strength sensor to WiZ (#67411)

---
 homeassistant/components/wiz/__init__.py |  8 ++-
 homeassistant/components/wiz/sensor.py   | 65 ++++++++++++++++++++++++
 tests/components/wiz/__init__.py         |  2 +-
 tests/components/wiz/test_sensor.py      | 37 ++++++++++++++
 4 files changed, 110 insertions(+), 2 deletions(-)
 create mode 100644 homeassistant/components/wiz/sensor.py
 create mode 100644 tests/components/wiz/test_sensor.py

diff --git a/homeassistant/components/wiz/__init__.py b/homeassistant/components/wiz/__init__.py
index d739c571c8b..ef6ffeb58b5 100644
--- a/homeassistant/components/wiz/__init__.py
+++ b/homeassistant/components/wiz/__init__.py
@@ -30,7 +30,13 @@ from .models import WizData
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS = [Platform.BINARY_SENSOR, Platform.LIGHT, Platform.NUMBER, Platform.SWITCH]
+PLATFORMS = [
+    Platform.BINARY_SENSOR,
+    Platform.LIGHT,
+    Platform.NUMBER,
+    Platform.SENSOR,
+    Platform.SWITCH,
+]
 
 REQUEST_REFRESH_DELAY = 0.35
 
diff --git a/homeassistant/components/wiz/sensor.py b/homeassistant/components/wiz/sensor.py
new file mode 100644
index 00000000000..d16130883d5
--- /dev/null
+++ b/homeassistant/components/wiz/sensor.py
@@ -0,0 +1,65 @@
+"""Support for WiZ sensors."""
+from __future__ import annotations
+
+from homeassistant.components.sensor import (
+    SensorDeviceClass,
+    SensorEntity,
+    SensorEntityDescription,
+    SensorStateClass,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import SIGNAL_STRENGTH_DECIBELS_MILLIWATT
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity import EntityCategory
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .const import DOMAIN
+from .entity import WizEntity
+from .models import WizData
+
+SENSORS: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key="rssi",
+        name="Signal Strength",
+        entity_registry_enabled_default=False,
+        state_class=SensorStateClass.MEASUREMENT,
+        device_class=SensorDeviceClass.SIGNAL_STRENGTH,
+        entity_category=EntityCategory.DIAGNOSTIC,
+        native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
+    ),
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up the wiz sensor."""
+    wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities(
+        WizSensor(wiz_data, entry.title, description) for description in SENSORS
+    )
+
+
+class WizSensor(WizEntity, SensorEntity):
+    """Defines a WiZ sensor."""
+
+    entity_description: SensorEntityDescription
+
+    def __init__(
+        self, wiz_data: WizData, name: str, description: SensorEntityDescription
+    ) -> None:
+        """Initialize an WiZ sensor."""
+        super().__init__(wiz_data, name)
+        self.entity_description = description
+        self._attr_unique_id = f"{self._device.mac}_{description.key}"
+        self._attr_name = f"{name} {description.name}"
+        self._async_update_attrs()
+
+    @callback
+    def _async_update_attrs(self) -> None:
+        """Handle updating _attr values."""
+        self._attr_native_value = self._device.state.pilotResult.get(
+            self.entity_description.key
+        )
diff --git a/tests/components/wiz/__init__.py b/tests/components/wiz/__init__.py
index c4a31b0a394..62920662c6f 100644
--- a/tests/components/wiz/__init__.py
+++ b/tests/components/wiz/__init__.py
@@ -199,7 +199,7 @@ def _mocked_wizlight(device, extended_white_range, bulb_type) -> wizlight:
     bulb.turn_on = AsyncMock()
     bulb.turn_off = AsyncMock()
     bulb.updateState = AsyncMock(return_value=FAKE_STATE)
-    bulb.getSupportedScenes = AsyncMock(return_value=list(SCENES))
+    bulb.getSupportedScenes = AsyncMock(return_value=list(SCENES.values()))
     bulb.start_push = AsyncMock(side_effect=_save_setup_callback)
     bulb.async_close = AsyncMock()
     bulb.set_speed = AsyncMock()
diff --git a/tests/components/wiz/test_sensor.py b/tests/components/wiz/test_sensor.py
new file mode 100644
index 00000000000..37a6b04dad3
--- /dev/null
+++ b/tests/components/wiz/test_sensor.py
@@ -0,0 +1,37 @@
+"""Tests for the sensor platform."""
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+
+from . import (
+    FAKE_DUAL_HEAD_RGBWW_BULB,
+    FAKE_MAC,
+    _patch_discovery,
+    _patch_wizlight,
+    async_push_update,
+    async_setup_integration,
+)
+
+
+async def test_signal_strength(hass: HomeAssistant) -> None:
+    """Test signal strength."""
+    bulb, entry = await async_setup_integration(
+        hass, bulb_type=FAKE_DUAL_HEAD_RGBWW_BULB
+    )
+    entity_id = "sensor.mock_title_signal_strength"
+    entity_registry = er.async_get(hass)
+    reg_entry = entity_registry.async_get(entity_id)
+    assert reg_entry.unique_id == f"{FAKE_MAC}_rssi"
+    updated_entity = entity_registry.async_update_entity(
+        entity_id=entity_id, disabled_by=None
+    )
+    assert not updated_entity.disabled
+
+    with _patch_discovery(), _patch_wizlight(device=bulb):
+        await hass.config_entries.async_reload(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert hass.states.get(entity_id).state == "-55"
+
+    await async_push_update(hass, bulb, {"mac": FAKE_MAC, "rssi": -50})
+    assert hass.states.get(entity_id).state == "-50"