Get MyStrom device state before checking support (#96004)

* Get device state before checking support

* Add full default device response to test

* Add test mocks

* Add test coverage

homeassistant/components/mystrom/__init__.py

@@ -22,6 +22,24 @@ PLATFORMS_BULB = [Platform.LIGHT]
 _LOGGER = logging.getLogger(__name__)
 
 
+async def _async_get_device_state(
+    device: MyStromSwitch | MyStromBulb, ip_address: str
+) -> None:
+    try:
+        await device.get_state()
+    except MyStromConnectionError as err:
+        _LOGGER.error("No route to myStrom plug: %s", ip_address)
+        raise ConfigEntryNotReady() from err
+
+
+def _get_mystrom_bulb(host: str, mac: str) -> MyStromBulb:
+    return MyStromBulb(host, mac)
+
+
+def _get_mystrom_switch(host: str) -> MyStromSwitch:
+    return MyStromSwitch(host)
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up myStrom from a config entry."""
     host = entry.data[CONF_HOST]
@@ -34,12 +52,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     device_type = info["type"]
     if device_type in [101, 106, 107]:
-        device = MyStromSwitch(host)
+        device = _get_mystrom_switch(host)
         platforms = PLATFORMS_SWITCH
-    elif device_type == 102:
+        await _async_get_device_state(device, info["ip"])
+    elif device_type in [102, 105]:
         mac = info["mac"]
-        device = MyStromBulb(host, mac)
+        device = _get_mystrom_bulb(host, mac)
         platforms = PLATFORMS_BULB
+        await _async_get_device_state(device, info["ip"])
         if device.bulb_type not in ["rgblamp", "strip"]:
             _LOGGER.error(
                 "Device %s (%s) is not a myStrom bulb nor myStrom LED Strip",
@@ -51,12 +71,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         _LOGGER.error("Unsupported myStrom device type: %s", device_type)
         return False
 
-    try:
-        await device.get_state()
-    except MyStromConnectionError as err:
-        _LOGGER.error("No route to myStrom plug: %s", info["ip"])
-        raise ConfigEntryNotReady() from err
-
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = MyStromData(
         device=device,
         info=info,
@@ -69,10 +83,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     device_type = hass.data[DOMAIN][entry.entry_id].info["type"]
+    platforms = []
     if device_type in [101, 106, 107]:
-        platforms = PLATFORMS_SWITCH
+        platforms.extend(PLATFORMS_SWITCH)
-    elif device_type == 102:
+    elif device_type in [102, 105]:
-        platforms = PLATFORMS_BULB
+        platforms.extend(PLATFORMS_BULB)
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, platforms):
         hass.data[DOMAIN].pop(entry.entry_id)
 

tests/components/mystrom/__init__.py

@@ -1 +1,172 @@
 """Tests for the myStrom integration."""
+from typing import Any, Optional
+
+
+def get_default_device_response(device_type: int) -> dict[str, Any]:
+    """Return default device response."""
+    return {
+        "version": "2.59.32",
+        "mac": "6001940376EB",
+        "type": device_type,
+        "ssid": "personal",
+        "ip": "192.168.0.23",
+        "mask": "255.255.255.0",
+        "gw": "192.168.0.1",
+        "dns": "192.168.0.1",
+        "static": False,
+        "connected": True,
+        "signal": 94,
+    }
+
+
+def get_default_bulb_state() -> dict[str, Any]:
+    """Get default bulb state."""
+    return {
+        "type": "rgblamp",
+        "battery": False,
+        "reachable": True,
+        "meshroot": True,
+        "on": False,
+        "color": "46;18;100",
+        "mode": "hsv",
+        "ramp": 10,
+        "power": 0.45,
+        "fw_version": "2.58.0",
+    }
+
+
+def get_default_switch_state() -> dict[str, Any]:
+    """Get default switch state."""
+    return {
+        "power": 1.69,
+        "Ws": 0.81,
+        "relay": True,
+        "temperature": 24.87,
+        "version": "2.59.32",
+        "mac": "6001940376EB",
+        "ssid": "personal",
+        "ip": "192.168.0.23",
+        "mask": "255.255.255.0",
+        "gw": "192.168.0.1",
+        "dns": "192.168.0.1",
+        "static": False,
+        "connected": True,
+        "signal": 94,
+    }
+
+
+class MyStromDeviceMock:
+    """Base device mock."""
+
+    def __init__(self, state: dict[str, Any]) -> None:
+        """Initialize device mock."""
+        self._requested_state = False
+        self._state = state
+
+    async def get_state(self) -> None:
+        """Set if state is requested."""
+        self._requested_state = True
+
+
+class MyStromBulbMock(MyStromDeviceMock):
+    """MyStrom Bulb mock."""
+
+    def __init__(self, mac: str, state: dict[str, Any]) -> None:
+        """Initialize bulb mock."""
+        super().__init__(state)
+        self.mac = mac
+
+    @property
+    def firmware(self) -> Optional[str]:
+        """Return current firmware."""
+        if not self._requested_state:
+            return None
+        return self._state["fw_version"]
+
+    @property
+    def consumption(self) -> Optional[float]:
+        """Return current firmware."""
+        if not self._requested_state:
+            return None
+        return self._state["power"]
+
+    @property
+    def color(self) -> Optional[str]:
+        """Return current color settings."""
+        if not self._requested_state:
+            return None
+        return self._state["color"]
+
+    @property
+    def mode(self) -> Optional[str]:
+        """Return current mode."""
+        if not self._requested_state:
+            return None
+        return self._state["mode"]
+
+    @property
+    def transition_time(self) -> Optional[int]:
+        """Return current transition time (ramp)."""
+        if not self._requested_state:
+            return None
+        return self._state["ramp"]
+
+    @property
+    def bulb_type(self) -> Optional[str]:
+        """Return the type of the bulb."""
+        if not self._requested_state:
+            return None
+        return self._state["type"]
+
+    @property
+    def state(self) -> Optional[bool]:
+        """Return the current state of the bulb."""
+        if not self._requested_state:
+            return None
+        return self._state["on"]
+
+
+class MyStromSwitchMock(MyStromDeviceMock):
+    """MyStrom Switch mock."""
+
+    @property
+    def relay(self) -> Optional[bool]:
+        """Return the relay state."""
+        if not self._requested_state:
+            return None
+        return self._state["on"]
+
+    @property
+    def consumption(self) -> Optional[float]:
+        """Return the current power consumption in mWh."""
+        if not self._requested_state:
+            return None
+        return self._state["power"]
+
+    @property
+    def consumedWs(self) -> Optional[float]:
+        """The average of energy consumed per second since last report call."""
+        if not self._requested_state:
+            return None
+        return self._state["Ws"]
+
+    @property
+    def firmware(self) -> Optional[str]:
+        """Return the current firmware."""
+        if not self._requested_state:
+            return None
+        return self._state["version"]
+
+    @property
+    def mac(self) -> Optional[str]:
+        """Return the MAC address."""
+        if not self._requested_state:
+            return None
+        return self._state["mac"]
+
+    @property
+    def temperature(self) -> Optional[float]:
+        """Return the current temperature in celsius."""
+        if not self._requested_state:
+            return None
+        return self._state["temperature"]

tests/components/mystrom/test_init.py

@@ -2,11 +2,19 @@
 from unittest.mock import AsyncMock, PropertyMock, patch
 
 from pymystrom.exceptions import MyStromConnectionError
+import pytest
 
 from homeassistant.components.mystrom.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
+from . import (
+    MyStromBulbMock,
+    MyStromSwitchMock,
+    get_default_bulb_state,
+    get_default_device_response,
+    get_default_switch_state,
+)
 from .conftest import DEVICE_MAC
 
 from tests.common import MockConfigEntry
@@ -16,30 +24,21 @@ async def init_integration(
     hass: HomeAssistant,
     config_entry: MockConfigEntry,
     device_type: int,
-    bulb_type: str = "strip",
 ) -> None:
     """Inititialize integration for testing."""
     with patch(
         "pymystrom.get_device_info",
-        side_effect=AsyncMock(return_value={"type": device_type, "mac": DEVICE_MAC}),
+        side_effect=AsyncMock(return_value=get_default_device_response(device_type)),
-    ), patch("pymystrom.switch.MyStromSwitch.get_state", return_value={}), patch(
-        "pymystrom.bulb.MyStromBulb.get_state", return_value={}
     ), patch(
-        "pymystrom.bulb.MyStromBulb.bulb_type", bulb_type
+        "homeassistant.components.mystrom._get_mystrom_bulb",
+        return_value=MyStromBulbMock("6001940376EB", get_default_bulb_state()),
     ), patch(
-        "pymystrom.switch.MyStromSwitch.mac",
+        "homeassistant.components.mystrom._get_mystrom_switch",
-        new_callable=PropertyMock,
+        return_value=MyStromSwitchMock(get_default_switch_state()),
-        return_value=DEVICE_MAC,
-    ), patch(
-        "pymystrom.bulb.MyStromBulb.mac",
-        new_callable=PropertyMock,
-        return_value=DEVICE_MAC,
     ):
         await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
-    assert config_entry.state == ConfigEntryState.LOADED
-
 
 async def test_init_switch_and_unload(
     hass: HomeAssistant, config_entry: MockConfigEntry
@@ -56,12 +55,35 @@ async def test_init_switch_and_unload(
     assert not hass.data.get(DOMAIN)
 
 
-async def test_init_bulb(hass: HomeAssistant, config_entry: MockConfigEntry) -> None:
+@pytest.mark.parametrize(
+    ("device_type", "platform", "entry_state", "entity_state_none"),
+    [
+        (101, "switch", ConfigEntryState.LOADED, False),
+        (102, "light", ConfigEntryState.LOADED, False),
+        (103, "button", ConfigEntryState.SETUP_ERROR, True),
+        (104, "button", ConfigEntryState.SETUP_ERROR, True),
+        (105, "light", ConfigEntryState.LOADED, False),
+        (106, "switch", ConfigEntryState.LOADED, False),
+        (107, "switch", ConfigEntryState.LOADED, False),
+        (110, "sensor", ConfigEntryState.SETUP_ERROR, True),
+        (113, "switch", ConfigEntryState.SETUP_ERROR, True),
+        (118, "button", ConfigEntryState.SETUP_ERROR, True),
+        (120, "switch", ConfigEntryState.SETUP_ERROR, True),
+    ],
+)
+async def test_init_bulb(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    device_type: int,
+    platform: str,
+    entry_state: ConfigEntryState,
+    entity_state_none: bool,
+) -> None:
     """Test the initialization of a myStrom bulb."""
-    await init_integration(hass, config_entry, 102)
+    await init_integration(hass, config_entry, device_type)
-    state = hass.states.get("light.mystrom_device")
+    state = hass.states.get(f"{platform}.mystrom_device")
-    assert state is not None
+    assert (state is None) == entity_state_none
-    assert config_entry.state is ConfigEntryState.LOADED
+    assert config_entry.state is entry_state
 
 
 async def test_init_of_unknown_bulb(
@@ -120,7 +142,7 @@ async def test_init_cannot_connect_because_of_get_state(
     """Test error handling for failing get_state."""
     with patch(
         "pymystrom.get_device_info",
-        side_effect=AsyncMock(return_value={"type": 101, "mac": DEVICE_MAC}),
+        side_effect=AsyncMock(return_value=get_default_device_response(101)),
     ), patch(
         "pymystrom.switch.MyStromSwitch.get_state", side_effect=MyStromConnectionError()
     ), patch(
@@ -129,4 +151,4 @@ async def test_init_cannot_connect_because_of_get_state(
         await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
 
-    assert config_entry.state == ConfigEntryState.SETUP_ERROR
+    assert config_entry.state == ConfigEntryState.SETUP_RETRY