From 901011de7b2f3c47ef43d0803c8d9555c51a16b9 Mon Sep 17 00:00:00 2001
From: Simone Chemelli <simone.chemelli@gmail.com>
Date: Wed, 19 Feb 2025 22:47:23 +0100
Subject: [PATCH] Use xmod model info for Shelly XMOD devices (#137013)

---
 .../components/shelly/config_flow.py          | 10 +++++++---
 .../components/shelly/coordinator.py          |  5 +++--
 homeassistant/components/shelly/utils.py      |  8 ++++++++
 tests/components/shelly/conftest.py           |  1 +
 tests/components/shelly/test_coordinator.py   | 19 +++++++++++++++++++
 5 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/shelly/config_flow.py b/homeassistant/components/shelly/config_flow.py
index 45655745403..5c5e187a0f4 100644
--- a/homeassistant/components/shelly/config_flow.py
+++ b/homeassistant/components/shelly/config_flow.py
@@ -7,7 +7,12 @@ from typing import Any, Final
 
 from aioshelly.block_device import BlockDevice
 from aioshelly.common import ConnectionOptions, get_info
-from aioshelly.const import BLOCK_GENERATIONS, DEFAULT_HTTP_PORT, RPC_GENERATIONS
+from aioshelly.const import (
+    BLOCK_GENERATIONS,
+    DEFAULT_HTTP_PORT,
+    MODEL_WALL_DISPLAY,
+    RPC_GENERATIONS,
+)
 from aioshelly.exceptions import (
     CustomPortNotSupported,
     DeviceConnectionError,
@@ -41,7 +46,6 @@ from .const import (
     CONF_SLEEP_PERIOD,
     DOMAIN,
     LOGGER,
-    MODEL_WALL_DISPLAY,
     BLEScannerMode,
 )
 from .coordinator import async_reconnect_soon
@@ -112,7 +116,7 @@ async def validate_input(
         return {
             "title": rpc_device.name,
             CONF_SLEEP_PERIOD: sleep_period,
-            "model": rpc_device.shelly.get("model"),
+            "model": rpc_device.xmod_info.get("p") or rpc_device.shelly.get("model"),
             CONF_GEN: gen,
         }
 
diff --git a/homeassistant/components/shelly/coordinator.py b/homeassistant/components/shelly/coordinator.py
index ad35ec32299..23d5842f4e4 100644
--- a/homeassistant/components/shelly/coordinator.py
+++ b/homeassistant/components/shelly/coordinator.py
@@ -10,7 +10,7 @@ from typing import Any, cast
 
 from aioshelly.ble import async_ensure_ble_enabled, async_stop_scanner
 from aioshelly.block_device import BlockDevice, BlockUpdateType
-from aioshelly.const import MODEL_NAMES, MODEL_VALVE
+from aioshelly.const import MODEL_VALVE
 from aioshelly.exceptions import (
     DeviceConnectionError,
     InvalidAuthError,
@@ -68,6 +68,7 @@ from .utils import (
     async_create_issue_unsupported_firmware,
     get_block_device_sleep_period,
     get_device_entry_gen,
+    get_device_info_model,
     get_host,
     get_http_port,
     get_rpc_device_wakeup_period,
@@ -164,7 +165,7 @@ class ShellyCoordinatorBase[_DeviceT: BlockDevice | RpcDevice](
             connections={(CONNECTION_NETWORK_MAC, self.mac)},
             identifiers={(DOMAIN, self.mac)},
             manufacturer="Shelly",
-            model=MODEL_NAMES.get(self.model),
+            model=get_device_info_model(self.device),
             model_id=self.model,
             sw_version=self.sw_version,
             hw_version=f"gen{get_device_entry_gen(self.config_entry)}",
diff --git a/homeassistant/components/shelly/utils.py b/homeassistant/components/shelly/utils.py
index fa310104424..4d3add7b17b 100644
--- a/homeassistant/components/shelly/utils.py
+++ b/homeassistant/components/shelly/utils.py
@@ -315,6 +315,14 @@ def get_model_name(info: dict[str, Any]) -> str:
     return cast(str, MODEL_NAMES.get(info["type"], info["type"]))
 
 
+def get_device_info_model(device: BlockDevice | RpcDevice) -> str | None:
+    """Return the device model for deviceinfo."""
+    if isinstance(device, RpcDevice) and (model := device.xmod_info.get("n")):
+        return cast(str, model)
+
+    return cast(str, MODEL_NAMES.get(device.model))
+
+
 def get_rpc_channel_name(device: RpcDevice, key: str) -> str:
     """Get name based on device and channel name."""
     key = key.replace("emdata", "em")
diff --git a/tests/components/shelly/conftest.py b/tests/components/shelly/conftest.py
index b3074742949..56b21701efe 100644
--- a/tests/components/shelly/conftest.py
+++ b/tests/components/shelly/conftest.py
@@ -476,6 +476,7 @@ def _mock_rpc_device(version: str | None = None):
         script_getcode=AsyncMock(
             side_effect=lambda script_id: {"data": MOCK_SCRIPTS[script_id - 1]}
         ),
+        xmod_info={},
     )
     type(device).name = PropertyMock(return_value="Test name")
     return device
diff --git a/tests/components/shelly/test_coordinator.py b/tests/components/shelly/test_coordinator.py
index 090c5e7207f..8c011e4ad0d 100644
--- a/tests/components/shelly/test_coordinator.py
+++ b/tests/components/shelly/test_coordinator.py
@@ -1011,3 +1011,22 @@ async def test_rpc_already_connected(
 
     assert "already connected" in caplog.text
     mock_rpc_device.initialize.assert_called_once()
+
+
+async def test_xmod_model_lookup(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test XMOD model look-up."""
+    xmod_model = "Test XMOD model name"
+    monkeypatch.setattr(mock_rpc_device, "xmod_info", {"n": xmod_model})
+    entry = await init_integration(hass, 2)
+
+    device = device_registry.async_get_device(
+        identifiers={(DOMAIN, entry.entry_id)},
+        connections={(dr.CONNECTION_NETWORK_MAC, dr.format_mac(entry.unique_id))},
+    )
+    assert device
+    assert device.model == xmod_model