Add support for Shelly Gen3 devices (#104874)

* Add support for Gen3 devices

* Add RPC_GENERATIONS const

* Add gen3 to tests

* More tests

* Add BLOCK_GENERATIONS const

* Use *_GENERATIONS constants from aioshelly

homeassistant/components/shelly/__init__.py

@@ -6,6 +6,7 @@
 
 from aioshelly.block_device import BlockDevice, BlockUpdateType
 from aioshelly.common import ConnectionOptions
+from aioshelly.const import RPC_GENERATIONS
 from aioshelly.exceptions import (
     DeviceConnectionError,
     InvalidAuthError,
@@ -123,7 +124,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     get_entry_data(hass)[entry.entry_id] = ShellyEntryData()
 
-    if get_device_entry_gen(entry) == 2:
+    if get_device_entry_gen(entry) in RPC_GENERATIONS:
         return await _async_setup_rpc_entry(hass, entry)
 
     return await _async_setup_block_entry(hass, entry)
@@ -313,7 +314,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if not entry.data.get(CONF_SLEEP_PERIOD):
         platforms = RPC_PLATFORMS
 
-    if get_device_entry_gen(entry) == 2:
+    if get_device_entry_gen(entry) in RPC_GENERATIONS:
         if unload_ok := await hass.config_entries.async_unload_platforms(
             entry, platforms
         ):

homeassistant/components/shelly/binary_sensor.py

@@ -4,6 +4,8 @@
 from dataclasses import dataclass
 from typing import Final, cast
 
+from aioshelly.const import RPC_GENERATIONS
+
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
@@ -224,7 +226,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up sensors for device."""
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         if config_entry.data[CONF_SLEEP_PERIOD]:
             async_setup_entry_rpc(
                 hass,

homeassistant/components/shelly/button.py

@@ -5,6 +5,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar
 
+from aioshelly.const import RPC_GENERATIONS
+
 from homeassistant.components.button import (
     ButtonDeviceClass,
     ButtonEntity,
@@ -126,7 +128,7 @@ def _async_migrate_unique_ids(
         return async_migrate_unique_ids(entity_entry, coordinator)
 
     coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None = None
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         coordinator = get_entry_data(hass)[config_entry.entry_id].rpc
     else:
         coordinator = get_entry_data(hass)[config_entry.entry_id].block

homeassistant/components/shelly/climate.py

@@ -6,6 +6,7 @@
 from typing import Any, cast
 
 from aioshelly.block_device import Block
+from aioshelly.const import RPC_GENERATIONS
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
 
 from homeassistant.components.climate import (
@@ -51,7 +52,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up climate device."""
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         return async_setup_rpc_entry(hass, config_entry, async_add_entities)
 
     coordinator = get_entry_data(hass)[config_entry.entry_id].block

homeassistant/components/shelly/config_flow.py

@@ -6,6 +6,7 @@
 
 from aioshelly.block_device import BlockDevice
 from aioshelly.common import ConnectionOptions, get_info
+from aioshelly.const import BLOCK_GENERATIONS, RPC_GENERATIONS
 from aioshelly.exceptions import (
     DeviceConnectionError,
     FirmwareUnsupported,
@@ -66,7 +67,9 @@ async def validate_input(
     """
     options = ConnectionOptions(host, data.get(CONF_USERNAME), data.get(CONF_PASSWORD))
 
-    if get_info_gen(info) == 2:
+    gen = get_info_gen(info)
+
+    if gen in RPC_GENERATIONS:
         ws_context = await get_ws_context(hass)
         rpc_device = await RpcDevice.create(
             async_get_clientsession(hass),
@@ -81,7 +84,7 @@ async def validate_input(
             "title": rpc_device.name,
             CONF_SLEEP_PERIOD: sleep_period,
             "model": rpc_device.shelly.get("model"),
-            "gen": 2,
+            "gen": gen,
         }
 
     # Gen1
@@ -96,7 +99,7 @@ async def validate_input(
         "title": block_device.name,
         CONF_SLEEP_PERIOD: get_block_device_sleep_period(block_device.settings),
         "model": block_device.model,
-        "gen": 1,
+        "gen": gen,
     }
 
 
@@ -165,7 +168,7 @@ async def async_step_credentials(
         """Handle the credentials step."""
         errors: dict[str, str] = {}
         if user_input is not None:
-            if get_info_gen(self.info) == 2:
+            if get_info_gen(self.info) in RPC_GENERATIONS:
                 user_input[CONF_USERNAME] = "admin"
             try:
                 device_info = await validate_input(
@@ -194,7 +197,7 @@ async def async_step_credentials(
         else:
             user_input = {}
 
-        if get_info_gen(self.info) == 2:
+        if get_info_gen(self.info) in RPC_GENERATIONS:
             schema = {
                 vol.Required(CONF_PASSWORD, default=user_input.get(CONF_PASSWORD)): str,
             }
@@ -331,7 +334,7 @@ async def async_step_reauth_confirm(
             await self.hass.config_entries.async_reload(self.entry.entry_id)
             return self.async_abort(reason="reauth_successful")
 
-        if self.entry.data.get("gen", 1) == 1:
+        if self.entry.data.get("gen", 1) in BLOCK_GENERATIONS:
             schema = {
                 vol.Required(CONF_USERNAME): str,
                 vol.Required(CONF_PASSWORD): str,
@@ -360,7 +363,7 @@ def async_get_options_flow(config_entry: ConfigEntry) -> OptionsFlowHandler:
     def async_supports_options_flow(cls, config_entry: ConfigEntry) -> bool:
         """Return options flow support for this handler."""
         return (
-            config_entry.data.get("gen") == 2
+            config_entry.data.get("gen") in RPC_GENERATIONS
             and not config_entry.data.get(CONF_SLEEP_PERIOD)
             and config_entry.data.get("model") != MODEL_WALL_DISPLAY
         )

homeassistant/components/shelly/cover.py

@@ -4,6 +4,7 @@
 from typing import Any, cast
 
 from aioshelly.block_device import Block
+from aioshelly.const import RPC_GENERATIONS
 
 from homeassistant.components.cover import (
     ATTR_POSITION,
@@ -26,7 +27,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up covers for device."""
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         return async_setup_rpc_entry(hass, config_entry, async_add_entities)
 
     return async_setup_block_entry(hass, config_entry, async_add_entities)

homeassistant/components/shelly/event.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, Any, Final
 
 from aioshelly.block_device import Block
-from aioshelly.const import MODEL_I3
+from aioshelly.const import MODEL_I3, RPC_GENERATIONS
 
 from homeassistant.components.event import (
     DOMAIN as EVENT_DOMAIN,
@@ -80,7 +80,7 @@ async def async_setup_entry(
 
     coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None = None
 
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         coordinator = get_entry_data(hass)[config_entry.entry_id].rpc
         if TYPE_CHECKING:
             assert coordinator

homeassistant/components/shelly/light.py

@@ -4,7 +4,7 @@
 from typing import Any, cast
 
 from aioshelly.block_device import Block
-from aioshelly.const import MODEL_BULB
+from aioshelly.const import MODEL_BULB, RPC_GENERATIONS
 
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
@@ -53,7 +53,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up lights for device."""
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         return async_setup_rpc_entry(hass, config_entry, async_add_entities)
 
     return async_setup_block_entry(hass, config_entry, async_add_entities)

homeassistant/components/shelly/sensor.py

@@ -6,6 +6,7 @@
 from typing import Final, cast
 
 from aioshelly.block_device import Block
+from aioshelly.const import RPC_GENERATIONS
 
 from homeassistant.components.sensor import (
     RestoreSensor,
@@ -925,7 +926,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up sensors for device."""
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         if config_entry.data[CONF_SLEEP_PERIOD]:
             async_setup_entry_rpc(
                 hass,

homeassistant/components/shelly/switch.py

@@ -5,7 +5,7 @@
 from typing import Any, cast
 
 from aioshelly.block_device import Block
-from aioshelly.const import MODEL_2, MODEL_25, MODEL_GAS
+from aioshelly.const import MODEL_2, MODEL_25, MODEL_GAS, RPC_GENERATIONS
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.config_entries import ConfigEntry
@@ -49,7 +49,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up switches for device."""
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         return async_setup_rpc_entry(hass, config_entry, async_add_entities)
 
     return async_setup_block_entry(hass, config_entry, async_add_entities)

homeassistant/components/shelly/update.py

@@ -6,6 +6,7 @@
 import logging
 from typing import Any, Final, cast
 
+from aioshelly.const import RPC_GENERATIONS
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
 
 from homeassistant.components.update import (
@@ -119,7 +120,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up update entities for Shelly component."""
-    if get_device_entry_gen(config_entry) == 2:
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
         if config_entry.data[CONF_SLEEP_PERIOD]:
             async_setup_entry_rpc(
                 hass,

homeassistant/components/shelly/utils.py

@@ -7,12 +7,14 @@
 from aiohttp.web import Request, WebSocketResponse
 from aioshelly.block_device import COAP, Block, BlockDevice
 from aioshelly.const import (
+    BLOCK_GENERATIONS,
     MODEL_1L,
     MODEL_DIMMER,
     MODEL_DIMMER_2,
     MODEL_EM3,
     MODEL_I3,
     MODEL_NAMES,
+    RPC_GENERATIONS,
 )
 from aioshelly.rpc_device import RpcDevice, WsServer
 
@@ -284,7 +286,7 @@ def get_info_gen(info: dict[str, Any]) -> int:
 
 def get_model_name(info: dict[str, Any]) -> str:
     """Return the device model name."""
-    if get_info_gen(info) == 2:
+    if get_info_gen(info) in RPC_GENERATIONS:
         return cast(str, MODEL_NAMES.get(info["model"], info["model"]))
 
     return cast(str, MODEL_NAMES.get(info["type"], info["type"]))
@@ -420,4 +422,4 @@ def get_release_url(gen: int, model: str, beta: bool) -> str | None:
     if beta or model in DEVICES_WITHOUT_FIRMWARE_CHANGELOG:
         return None
 
-    return GEN1_RELEASE_URL if gen == 1 else GEN2_RELEASE_URL
+    return GEN1_RELEASE_URL if gen in BLOCK_GENERATIONS else GEN2_RELEASE_URL

tests/components/shelly/test_config_flow.py

@@ -55,6 +55,7 @@
     [
         (1, MODEL_1),
         (2, MODEL_PLUS_2PM),
+        (3, MODEL_PLUS_2PM),
     ],
 )
 async def test_form(
@@ -109,6 +110,12 @@ async def test_form(
             {"password": "test2 password"},
             "admin",
         ),
+        (
+            3,
+            MODEL_PLUS_2PM,
+            {"password": "test2 password"},
+            "admin",
+        ),
     ],
 )
 async def test_form_auth(
@@ -465,6 +472,11 @@ async def test_form_auth_errors_test_connection_gen2(
             MODEL_PLUS_2PM,
             {"mac": "test-mac", "model": MODEL_PLUS_2PM, "auth": False, "gen": 2},
         ),
+        (
+            3,
+            MODEL_PLUS_2PM,
+            {"mac": "test-mac", "model": MODEL_PLUS_2PM, "auth": False, "gen": 3},
+        ),
     ],
 )
 async def test_zeroconf(
@@ -742,6 +754,7 @@ async def test_zeroconf_require_auth(hass: HomeAssistant, mock_block_device) ->
     [
         (1, {"username": "test user", "password": "test1 password"}),
         (2, {"password": "test2 password"}),
+        (3, {"password": "test2 password"}),
     ],
 )
 async def test_reauth_successful(
@@ -780,6 +793,7 @@ async def test_reauth_successful(
     [
         (1, {"username": "test user", "password": "test1 password"}),
         (2, {"password": "test2 password"}),
+        (3, {"password": "test2 password"}),
     ],
 )
 async def test_reauth_unsuccessful(hass: HomeAssistant, gen, user_input) -> None:

tests/components/shelly/test_init.py

@@ -41,7 +41,7 @@ async def test_custom_coap_port(
     assert "Starting CoAP context with UDP port 7632" in caplog.text
 
 
-@pytest.mark.parametrize("gen", [1, 2])
+@pytest.mark.parametrize("gen", [1, 2, 3])
 async def test_shared_device_mac(
     hass: HomeAssistant,
     gen,
@@ -74,7 +74,7 @@ async def test_setup_entry_not_shelly(
     assert "probably comes from a custom integration" in caplog.text
 
 
-@pytest.mark.parametrize("gen", [1, 2])
+@pytest.mark.parametrize("gen", [1, 2, 3])
 async def test_device_connection_error(
     hass: HomeAssistant, gen, mock_block_device, mock_rpc_device, monkeypatch
 ) -> None:
@@ -90,7 +90,7 @@ async def test_device_connection_error(
     assert entry.state == ConfigEntryState.SETUP_RETRY
 
 
-@pytest.mark.parametrize("gen", [1, 2])
+@pytest.mark.parametrize("gen", [1, 2, 3])
 async def test_mac_mismatch_error(
     hass: HomeAssistant, gen, mock_block_device, mock_rpc_device, monkeypatch
 ) -> None:
@@ -106,7 +106,7 @@ async def test_mac_mismatch_error(
     assert entry.state == ConfigEntryState.SETUP_RETRY
 
 
-@pytest.mark.parametrize("gen", [1, 2])
+@pytest.mark.parametrize("gen", [1, 2, 3])
 async def test_device_auth_error(
     hass: HomeAssistant, gen, mock_block_device, mock_rpc_device, monkeypatch
 ) -> None: