Fix Shelly model name for xmod devices (#138984)

homeassistant/components/shelly/coordinator.py

@@ -68,11 +68,11 @@ 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,
     get_rpc_ws_url,
+    get_shelly_model_name,
     update_device_fw_info,
 )
 
@@ -165,7 +165,7 @@ class ShellyCoordinatorBase[_DeviceT: BlockDevice | RpcDevice](
             connections={(CONNECTION_NETWORK_MAC, self.mac)},
             identifiers={(DOMAIN, self.mac)},
             manufacturer="Shelly",
-            model=get_device_info_model(self.device),
+            model=get_shelly_model_name(self.model, self.sleep_period, self.device),
             model_id=self.model,
             sw_version=self.sw_version,
             hw_version=f"gen{get_device_entry_gen(self.config_entry)}",

homeassistant/components/shelly/utils.py

@@ -315,12 +315,25 @@ 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)
+def get_shelly_model_name(
+    model: str,
+    sleep_period: int,
+    device: BlockDevice | RpcDevice,
+) -> str | None:
+    """Get Shelly model name.
 
-    return cast(str, MODEL_NAMES.get(device.model))
+    Assume that XMOD devices are not sleepy devices.
+    """
+    if (
+        sleep_period == 0
+        and isinstance(device, RpcDevice)
+        and (model_name := device.xmod_info.get("n"))
+    ):
+        # Use the model name from XMOD data
+        return cast(str, model_name)
+
+    # Use the model name from aioshelly
+    return cast(str, MODEL_NAMES.get(model))
 
 
 def get_rpc_channel_name(device: RpcDevice, key: str) -> str:

tests/components/shelly/conftest.py

@@ -13,6 +13,7 @@ from aioshelly.ble.const import (
 )
 from aioshelly.block_device import BlockDevice, BlockUpdateType
 from aioshelly.const import MODEL_1, MODEL_25, MODEL_PLUS_2PM
+from aioshelly.exceptions import NotInitialized
 from aioshelly.rpc_device import RpcDevice, RpcUpdateType
 import pytest
 
@@ -568,3 +569,97 @@ async def mock_blu_trv():
         blu_trv_device_mock.return_value.mock_update = Mock(side_effect=update)
 
         yield blu_trv_device_mock.return_value
+
+
+def _mock_sleepy_not_initialized_rpc_device():
+    """Mock sleepy NotInitialized rpc (Gen2+, Websocket) device."""
+    device = Mock(spec=RpcDevice, initialized=False, connected=False)
+    type(device).requires_auth = PropertyMock(side_effect=NotInitialized)
+    type(device).status = PropertyMock(side_effect=NotInitialized)
+    type(device).event = PropertyMock(side_effect=NotInitialized)
+    type(device).config = PropertyMock(side_effect=NotInitialized)
+    type(device).shelly = PropertyMock(side_effect=NotInitialized)
+    type(device).gen = PropertyMock(side_effect=NotInitialized)
+    type(device).firmware_version = PropertyMock(side_effect=NotInitialized)
+    type(device).version = PropertyMock(side_effect=NotInitialized)
+    type(device).model = PropertyMock(side_effect=NotInitialized)
+    type(device).xmod_info = PropertyMock(side_effect=NotInitialized)
+    type(device).hostname = PropertyMock(side_effect=NotInitialized)
+    type(device).name = PropertyMock(side_effect=NotInitialized)
+    type(device).firmware_supported = PropertyMock(side_effect=NotInitialized)
+    return device
+
+
+def initialize_sleepy_rpc_device(device):
+    """Initialize a sleepy RPC (Gen2+, Websocket) device."""
+    type(device).requires_auth = PropertyMock()
+    type(device).status = PropertyMock(return_value=MOCK_STATUS_RPC)
+    type(device).event = PropertyMock(return_value={})
+    type(device).config = PropertyMock(return_value=MOCK_CONFIG)
+    type(device).shelly = PropertyMock(return_value=MOCK_SHELLY_RPC)
+    type(device).gen = PropertyMock(return_value=2)
+    type(device).firmware_version = PropertyMock(
+        return_value="20240425-141520/1.3.0-ga3fdd3d"
+    )
+    type(device).version = PropertyMock("1.3.0")
+    type(device).model = PropertyMock("SPSW-201PE16EU")
+    type(device).xmod_info = PropertyMock(return_value={})
+    type(device).hostname = PropertyMock(return_value="hostname")
+    type(device).name = PropertyMock(return_value="Test Name")
+    type(device).firmware_supported = PropertyMock(return_value=True)
+
+    device.status["sys"]["wakeup_period"] = 1000
+    device.connected = True
+    device.initialized = True
+
+
+@pytest.fixture
+async def mock_sleepy_rpc_device():
+    """Mock sleepy RPC (Gen2+, Websocket) device.
+
+    Mock a RPC device that is not initialized and raises NotInitialized
+    when aioshelly properties are accessed.
+
+    Initialize the device when initialize() method is called.
+    """
+    with patch("aioshelly.rpc_device.RpcDevice.create") as rpc_device_mock:
+
+        def update():
+            rpc_device_mock.return_value.subscribe_updates.call_args[0][0](
+                {}, RpcUpdateType.STATUS
+            )
+
+        def event():
+            rpc_device_mock.return_value.subscribe_updates.call_args[0][0](
+                {}, RpcUpdateType.EVENT
+            )
+
+        def online():
+            rpc_device_mock.return_value.subscribe_updates.call_args[0][0](
+                {}, RpcUpdateType.ONLINE
+            )
+
+        def disconnected():
+            rpc_device_mock.return_value.subscribe_updates.call_args[0][0](
+                {}, RpcUpdateType.DISCONNECTED
+            )
+
+        def initialized():
+            rpc_device_mock.return_value.subscribe_updates.call_args[0][0](
+                {}, RpcUpdateType.INITIALIZED
+            )
+
+        def _initialize():
+            initialize_sleepy_rpc_device(device)
+
+        device = _mock_sleepy_not_initialized_rpc_device()
+        device.initialize = AsyncMock(side_effect=_initialize)
+        rpc_device_mock.return_value = device
+
+        rpc_device_mock.return_value.mock_disconnected = Mock(side_effect=disconnected)
+        rpc_device_mock.return_value.mock_update = Mock(side_effect=update)
+        rpc_device_mock.return_value.mock_event = Mock(side_effect=event)
+        rpc_device_mock.return_value.mock_online = Mock(side_effect=online)
+        rpc_device_mock.return_value.mock_initialized = Mock(side_effect=initialized)
+
+        yield rpc_device_mock.return_value

tests/components/shelly/test_init.py

@@ -312,13 +312,10 @@ async def test_sleeping_rpc_device_online_new_firmware(
 
 async def test_sleeping_rpc_device_online_during_setup(
     hass: HomeAssistant,
-    mock_rpc_device: Mock,
-    monkeypatch: pytest.MonkeyPatch,
+    mock_sleepy_rpc_device: Mock,
     caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test sleeping device Gen2 woke up by user during setup."""
-    monkeypatch.setattr(mock_rpc_device, "connected", False)
-    monkeypatch.setitem(mock_rpc_device.status["sys"], "wakeup_period", 1000)
     await init_integration(hass, 2, sleep_period=1000)
     await hass.async_block_till_done(wait_background_tasks=True)