From dbe98de82a28769b0a7e66a3cc4366262d9c9044 Mon Sep 17 00:00:00 2001
From: Maciej Bieniek <bieniu@users.noreply.github.com>
Date: Fri, 5 Jul 2024 09:40:43 +0200
Subject: [PATCH] Fix `pulse counter frequency` sensors for Shelly Plus Uni
 (#121178)

Co-authored-by: Maciej Bieniek <478555+bieniu@users.noreply.github.com>
---
 homeassistant/components/shelly/sensor.py | 29 +++++++++++++++--------
 tests/components/shelly/conftest.py       |  4 +++-
 tests/components/shelly/test_sensor.py    | 26 ++++++++++++++++++++
 3 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/shelly/sensor.py b/homeassistant/components/shelly/sensor.py
index 743c7c7ff01..5a6f03fd90c 100644
--- a/homeassistant/components/shelly/sensor.py
+++ b/homeassistant/components/shelly/sensor.py
@@ -960,14 +960,18 @@ RPC_SENSORS: Final = {
         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,7 +981,9 @@ RPC_SENSORS: Final = {
         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",
@@ -985,26 +991,29 @@ RPC_SENSORS: Final = {
         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
index 65ebdeb6996..e9662b55cf5 100644
--- a/tests/components/shelly/conftest.py
+++ b/tests/components/shelly/conftest.py
@@ -228,7 +228,9 @@ MOCK_STATUS_RPC = {
     "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
index 513bcd875e2..c62a1f6f6ca 100644
--- 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