diff --git a/homeassistant/components/shelly/__init__.py b/homeassistant/components/shelly/__init__.py index 553d32f8e48..0fab86f7f4f 100644 --- a/homeassistant/components/shelly/__init__.py +++ b/homeassistant/components/shelly/__init__.py @@ -6,6 +6,7 @@ from typing import Any, Final 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 ): diff --git a/homeassistant/components/shelly/binary_sensor.py b/homeassistant/components/shelly/binary_sensor.py index a5889cd11a7..caed52279da 100644 --- a/homeassistant/components/shelly/binary_sensor.py +++ b/homeassistant/components/shelly/binary_sensor.py @@ -4,6 +4,8 @@ from __future__ import annotations 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, diff --git a/homeassistant/components/shelly/button.py b/homeassistant/components/shelly/button.py index edc33c9a8a0..e5cc6b6580b 100644 --- a/homeassistant/components/shelly/button.py +++ b/homeassistant/components/shelly/button.py @@ -5,6 +5,8 @@ from collections.abc import Callable, Coroutine 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 @@ async def async_setup_entry( 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 diff --git a/homeassistant/components/shelly/climate.py b/homeassistant/components/shelly/climate.py index 6a592c904f6..9ac603a7fb0 100644 --- a/homeassistant/components/shelly/climate.py +++ b/homeassistant/components/shelly/climate.py @@ -6,6 +6,7 @@ from dataclasses import asdict, dataclass 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 diff --git a/homeassistant/components/shelly/config_flow.py b/homeassistant/components/shelly/config_flow.py index 98233d27b22..68b0f1f8ccc 100644 --- a/homeassistant/components/shelly/config_flow.py +++ b/homeassistant/components/shelly/config_flow.py @@ -6,6 +6,7 @@ from typing import Any, Final 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 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): """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 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): 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 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): 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 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN): 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 ) diff --git a/homeassistant/components/shelly/cover.py b/homeassistant/components/shelly/cover.py index 95f387f8f97..4390790c794 100644 --- a/homeassistant/components/shelly/cover.py +++ b/homeassistant/components/shelly/cover.py @@ -4,6 +4,7 @@ from __future__ import annotations 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) diff --git a/homeassistant/components/shelly/event.py b/homeassistant/components/shelly/event.py index af323c82a24..e93303d7191 100644 --- a/homeassistant/components/shelly/event.py +++ b/homeassistant/components/shelly/event.py @@ -6,7 +6,7 @@ from dataclasses import dataclass 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 diff --git a/homeassistant/components/shelly/light.py b/homeassistant/components/shelly/light.py index 2dfc5b497b1..7e49dc78e4d 100644 --- a/homeassistant/components/shelly/light.py +++ b/homeassistant/components/shelly/light.py @@ -4,7 +4,7 @@ from __future__ import annotations 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) diff --git a/homeassistant/components/shelly/sensor.py b/homeassistant/components/shelly/sensor.py index 99ccd9ab2ff..4518135214c 100644 --- a/homeassistant/components/shelly/sensor.py +++ b/homeassistant/components/shelly/sensor.py @@ -6,6 +6,7 @@ from dataclasses import dataclass 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, diff --git a/homeassistant/components/shelly/switch.py b/homeassistant/components/shelly/switch.py index 5a398182e4d..5ef39cd33af 100644 --- a/homeassistant/components/shelly/switch.py +++ b/homeassistant/components/shelly/switch.py @@ -5,7 +5,7 @@ from dataclasses import dataclass 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) diff --git a/homeassistant/components/shelly/update.py b/homeassistant/components/shelly/update.py index 9e52a292108..975b61e631a 100644 --- a/homeassistant/components/shelly/update.py +++ b/homeassistant/components/shelly/update.py @@ -6,6 +6,7 @@ from dataclasses import dataclass 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, diff --git a/homeassistant/components/shelly/utils.py b/homeassistant/components/shelly/utils.py index b53e3153a09..7d475bf5ef8 100644 --- a/homeassistant/components/shelly/utils.py +++ b/homeassistant/components/shelly/utils.py @@ -7,12 +7,14 @@ from typing import Any, cast 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 diff --git a/tests/components/shelly/test_config_flow.py b/tests/components/shelly/test_config_flow.py index c7ac472ada4..1bccd3570cf 100644 --- a/tests/components/shelly/test_config_flow.py +++ b/tests/components/shelly/test_config_flow.py @@ -55,6 +55,7 @@ DISCOVERY_INFO_WITH_MAC = zeroconf.ZeroconfServiceInfo( [ (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: diff --git a/tests/components/shelly/test_init.py b/tests/components/shelly/test_init.py index 2ead9cba198..8f6599b39e4 100644 --- a/tests/components/shelly/test_init.py +++ b/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: