Fix pulse counter frequency sensors for Shelly Plus Uni (#121178)

Co-authored-by: Maciej Bieniek <[email protected]>
home-assistant · Jul 5, 2024 · ee276af · ee276af
1 parent 0acd1dc
commit ee276af
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 11 deletions.
diff --git a/homeassistant/components/shelly/sensor.py b/homeassistant/components/shelly/sensor.py
@@ -960,14 +960,18 @@ class RestSensorDescription(RestEntityDescription, SensorEntityDescription):
         name="Analog input",
         native_unit_of_measurement=PERCENTAGE,
         state_class=SensorStateClass.MEASUREMENT,
-        removal_condition=lambda config, _status, key: (config[key]["enable"] is False),
+        removal_condition=lambda config, _, key: (
+            config[key]["type"] != "analog" or config[key]["enable"] is False
+        ),
     ),
     "analoginput_xpercent": RpcSensorDescription(
         key="input",
         sub_key="xpercent",
         name="Analog value",
         removal_condition=lambda config, status, key: (
-            config[key]["enable"] is False or status[key].get("xpercent") is None
+            config[key]["type"] != "analog"
+            or config[key]["enable"] is False
+            or status[key].get("xpercent") is None
         ),
     ),
     "pulse_counter": RpcSensorDescription(
@@ -977,34 +981,39 @@ class RestSensorDescription(RestEntityDescription, SensorEntityDescription):
         native_unit_of_measurement="pulse",
         state_class=SensorStateClass.TOTAL,
         value=lambda status, _: status["total"],
-        removal_condition=lambda config, _status, key: (config[key]["enable"] is False),
+        removal_condition=lambda config, _status, key: (
+            config[key]["type"] != "count" or config[key]["enable"] is False
+        ),
     ),
     "counter_value": RpcSensorDescription(
         key="input",
         sub_key="counts",
         name="Counter value",
         value=lambda status, _: status["xtotal"],
         removal_condition=lambda config, status, key: (
-            config[key]["enable"] is False
+            config[key]["type"] != "count"
+            or config[key]["enable"] is False
             or status[key]["counts"].get("xtotal") is None
         ),
     ),
     "counter_frequency": RpcSensorDescription(
         key="input",
-        sub_key="counts",
+        sub_key="freq",
         name="Pulse counter frequency",
         native_unit_of_measurement=UnitOfFrequency.HERTZ,
         state_class=SensorStateClass.MEASUREMENT,
-        value=lambda status, _: status["freq"],
-        removal_condition=lambda config, status, key: (config[key]["enable"] is False),
+        removal_condition=lambda config, _, key: (
+            config[key]["type"] != "count" or config[key]["enable"] is False
+        ),
     ),
     "counter_frequency_value": RpcSensorDescription(
         key="input",
-        sub_key="counts",
+        sub_key="xfreq",
         name="Pulse counter frequency value",
-        value=lambda status, _: status["xfreq"],
         removal_condition=lambda config, status, key: (
-            config[key]["enable"] is False or status[key]["counts"].get("xfreq") is None
+            config[key]["type"] != "count"
+            or config[key]["enable"] is False
+            or status[key].get("xfreq") is None
         ),
     ),
 }

diff --git a/tests/components/shelly/conftest.py b/tests/components/shelly/conftest.py
@@ -228,7 +228,9 @@ def mock_white_light_set_state(
     "input:1": {"id": 1, "percent": 89, "xpercent": 8.9},
     "input:2": {
         "id": 2,
-        "counts": {"total": 56174, "xtotal": 561.74, "freq": 208.00, "xfreq": 6.11},
+        "counts": {"total": 56174, "xtotal": 561.74},
+        "freq": 208.00,
+        "xfreq": 6.11,
     },
     "light:0": {"output": True, "brightness": 53.0},
     "light:1": {"output": True, "brightness": 53.0},

diff --git a/tests/components/shelly/test_sensor.py b/tests/components/shelly/test_sensor.py
@@ -828,3 +828,29 @@ async def test_rpc_pulse_counter_frequency_sensors(
     entry = entity_registry.async_get(entity_id)
     assert entry
     assert entry.unique_id == "123456789ABC-input:2-counter_frequency_value"
+
+
+async def test_rpc_disabled_xfreq(
+    hass: HomeAssistant,
+    mock_rpc_device: Mock,
+    entity_registry: EntityRegistry,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test RPC input with the xfreq sensor disabled."""
+    status = deepcopy(mock_rpc_device.status)
+    status["input:2"] = {
+        "id": 2,
+        "counts": {"total": 56174, "xtotal": 561.74},
+        "freq": 208.00,
+    }
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    await init_integration(hass, 2)
+
+    entity_id = f"{SENSOR_DOMAIN}.gas_pulse_counter_frequency_value"
+
+    state = hass.states.get(entity_id)
+    assert not state
+
+    entry = entity_registry.async_get(entity_id)
+    assert not entry