From 70321ed795086f385dd521f04156a26ff0948ffd Mon Sep 17 00:00:00 2001
From: Patrik Lindgren <21142447+ggravlingen@users.noreply.github.com>
Date: Thu, 27 Jan 2022 17:47:47 +0100
Subject: [PATCH] Add battery sensor for Tradfri blinds (#65067)

---
 homeassistant/components/tradfri/sensor.py |  5 ++-
 tests/components/tradfri/test_sensor.py    | 42 ++++++++++++++++------
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/homeassistant/components/tradfri/sensor.py b/homeassistant/components/tradfri/sensor.py
index 3d042fa6417..3f8fa96857d 100644
--- a/homeassistant/components/tradfri/sensor.py
+++ b/homeassistant/components/tradfri/sensor.py
@@ -28,7 +28,7 @@ async def async_setup_entry(
     api = coordinator_data[KEY_API]
 
     async_add_entities(
-        TradfriSensor(
+        TradfriBatterySensor(
             device_coordinator,
             api,
             gateway_id,
@@ -37,14 +37,13 @@ async def async_setup_entry(
         if (
             not device_coordinator.device.has_light_control
             and not device_coordinator.device.has_socket_control
-            and not device_coordinator.device.has_blind_control
             and not device_coordinator.device.has_signal_repeater_control
             and not device_coordinator.device.has_air_purifier_control
         )
     )
 
 
-class TradfriSensor(TradfriBaseEntity, SensorEntity):
+class TradfriBatterySensor(TradfriBaseEntity, SensorEntity):
     """The platform class required by Home Assistant."""
 
     _attr_device_class = SensorDeviceClass.BATTERY
diff --git a/tests/components/tradfri/test_sensor.py b/tests/components/tradfri/test_sensor.py
index 1e9ae718285..63d4b6f84e5 100644
--- a/tests/components/tradfri/test_sensor.py
+++ b/tests/components/tradfri/test_sensor.py
@@ -5,16 +5,13 @@ from unittest.mock import MagicMock, Mock
 from .common import setup_integration
 
 
-def mock_sensor(state_name: str, state_value: str, device_number=0):
+def mock_sensor(test_state: list, device_number=0):
     """Mock a tradfri sensor."""
     dev_info_mock = MagicMock()
     dev_info_mock.manufacturer = "manufacturer"
     dev_info_mock.model_number = "model"
     dev_info_mock.firmware_version = "1.2.3"
 
-    # Set state value, eg battery_level = 50
-    setattr(dev_info_mock, state_name, state_value)
-
     _mock_sensor = Mock(
         id=f"mock-sensor-id-{device_number}",
         reachable=True,
@@ -26,6 +23,11 @@ def mock_sensor(state_name: str, state_value: str, device_number=0):
         has_signal_repeater_control=False,
         has_air_purifier_control=False,
     )
+
+    # Set state value, eg battery_level = 50, or has_air_purifier_control=True
+    for state in test_state:
+        setattr(dev_info_mock, state["attribute"], state["value"])
+
     _mock_sensor.name = f"tradfri_sensor_{device_number}"
 
     return _mock_sensor
@@ -34,7 +36,7 @@ def mock_sensor(state_name: str, state_value: str, device_number=0):
 async def test_battery_sensor(hass, mock_gateway, mock_api_factory):
     """Test that a battery sensor is correctly added."""
     mock_gateway.mock_devices.append(
-        mock_sensor(state_name="battery_level", state_value=60)
+        mock_sensor(test_state=[{"attribute": "battery_level", "value": 60}])
     )
     await setup_integration(hass)
 
@@ -45,10 +47,27 @@ async def test_battery_sensor(hass, mock_gateway, mock_api_factory):
     assert sensor_1.attributes["device_class"] == "battery"
 
 
+async def test_cover_battery_sensor(hass, mock_gateway, mock_api_factory):
+    """Test that a battery sensor is correctly added for a cover (blind)."""
+    mock_gateway.mock_devices.append(
+        mock_sensor(
+            test_state=[
+                {"attribute": "battery_level", "value": 42, "has_blind_control": True}
+            ]
+        )
+    )
+    await setup_integration(hass)
+
+    sensor_1 = hass.states.get("sensor.tradfri_sensor_0")
+    assert sensor_1 is not None
+    assert sensor_1.state == "42"
+    assert sensor_1.attributes["unit_of_measurement"] == "%"
+    assert sensor_1.attributes["device_class"] == "battery"
+
+
 async def test_sensor_observed(hass, mock_gateway, mock_api_factory):
     """Test that sensors are correctly observed."""
-
-    sensor = mock_sensor(state_name="battery_level", state_value=60)
+    sensor = mock_sensor(test_state=[{"attribute": "battery_level", "value": 60}])
     mock_gateway.mock_devices.append(sensor)
     await setup_integration(hass)
     assert len(sensor.observe.mock_calls) > 0
@@ -56,11 +75,14 @@ async def test_sensor_observed(hass, mock_gateway, mock_api_factory):
 
 async def test_sensor_available(hass, mock_gateway, mock_api_factory):
     """Test sensor available property."""
-
-    sensor = mock_sensor(state_name="battery_level", state_value=60, device_number=1)
+    sensor = mock_sensor(
+        test_state=[{"attribute": "battery_level", "value": 60}], device_number=1
+    )
     sensor.reachable = True
 
-    sensor2 = mock_sensor(state_name="battery_level", state_value=60, device_number=2)
+    sensor2 = mock_sensor(
+        test_state=[{"attribute": "battery_level", "value": 60}], device_number=2
+    )
     sensor2.reachable = False
 
     mock_gateway.mock_devices.append(sensor)