Fix Shelly entry unload and add tests for init (#80760)

.coveragerc

@@ -1106,7 +1106,6 @@ omit =
     homeassistant/components/sesame/lock.py
     homeassistant/components/seven_segments/image_processing.py
     homeassistant/components/seventeentrack/sensor.py
-    homeassistant/components/shelly/__init__.py
     homeassistant/components/shelly/binary_sensor.py
     homeassistant/components/shelly/climate.py
     homeassistant/components/shelly/coordinator.py

homeassistant/components/shelly/__init__.py

@@ -143,6 +143,7 @@ async def _async_setup_block_entry(hass: HomeAssistant, entry: ConfigEntry) -> b
                 )
             },
         )
+    # https://github.com/home-assistant/core/pull/48076
     if device_entry and entry.entry_id not in device_entry.config_entries:
         device_entry = None
 
@@ -231,6 +232,7 @@ async def _async_setup_rpc_entry(hass: HomeAssistant, entry: ConfigEntry) -> boo
                 )
             },
         )
+    # https://github.com/home-assistant/core/pull/48076
     if device_entry and entry.entry_id not in device_entry.config_entries:
         device_entry = None
 
@@ -295,9 +297,13 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     shelly_entry_data = get_entry_data(hass)[entry.entry_id]
 
-    if shelly_entry_data.device is not None:
+    # If device is present, block/rpc coordinator is not setup yet
-        # If device is present, block/rpc coordinator is not setup yet
+    device = shelly_entry_data.device
-        shelly_entry_data.device.shutdown()
+    if isinstance(device, RpcDevice):
+        await device.shutdown()
+        return True
+    if isinstance(device, BlockDevice):
+        device.shutdown()
         return True
 
     platforms = RPC_SLEEPING_PLATFORMS

tests/components/shelly/__init__.py

@@ -5,19 +5,21 @@ from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry
 
+MOCK_MAC = "123456789ABC"
+
 
 async def init_integration(
-    hass: HomeAssistant, gen: int, model="SHSW-25"
+    hass: HomeAssistant, gen: int, model="SHSW-25", sleep_period=0
 ) -> MockConfigEntry:
     """Set up the Shelly integration in Home Assistant."""
     data = {
         CONF_HOST: "192.168.1.37",
-        CONF_SLEEP_PERIOD: 0,
+        CONF_SLEEP_PERIOD: sleep_period,
         "model": model,
         "gen": gen,
     }
 
-    entry = MockConfigEntry(domain=DOMAIN, data=data, unique_id=DOMAIN)
+    entry = MockConfigEntry(domain=DOMAIN, data=data, unique_id=MOCK_MAC)
     entry.add_to_hass(hass)
 
     await hass.config_entries.async_setup(entry.entry_id)

tests/components/shelly/conftest.py

@@ -1,6 +1,8 @@
 """Test configuration for Shelly."""
 from unittest.mock import AsyncMock, Mock, patch
 
+from aioshelly.block_device import BlockDevice
+from aioshelly.rpc_device import RpcDevice
 import pytest
 
 from homeassistant.components.shelly.const import (
@@ -8,13 +10,15 @@ from homeassistant.components.shelly.const import (
     REST_SENSORS_UPDATE_INTERVAL,
 )
 
+from . import MOCK_MAC
+
 from tests.common import async_capture_events, async_mock_service, mock_device_registry
 
 MOCK_SETTINGS = {
     "name": "Test name",
     "mode": "relay",
     "device": {
-        "mac": "test-mac",
+        "mac": MOCK_MAC,
         "hostname": "test-host",
         "type": "SHSW-25",
         "num_outputs": 2,
@@ -95,7 +99,7 @@ MOCK_CONFIG = {
 }
 
 MOCK_SHELLY_COAP = {
-    "mac": "test-mac",
+    "mac": MOCK_MAC,
     "auth": False,
     "fw": "20201124-092854/v1.9.0@57ac4ad8",
     "num_outputs": 2,
@@ -104,7 +108,7 @@ MOCK_SHELLY_COAP = {
 MOCK_SHELLY_RPC = {
     "name": "Test Gen2",
     "id": "shellyplus2pm-123456789abc",
-    "mac": "123456789ABC",
+    "mac": MOCK_MAC,
     "model": "SNSW-002P16EU",
     "gen": 2,
     "fw_id": "20220830-130540/0.11.0-gfa1bc37",
@@ -142,7 +146,13 @@ MOCK_STATUS_RPC = {
 @pytest.fixture(autouse=True)
 def mock_coap():
     """Mock out coap."""
-    with patch("homeassistant.components.shelly.utils.get_coap_context"):
+    with patch(
+        "homeassistant.components.shelly.utils.COAP",
+        return_value=Mock(
+            initialize=AsyncMock(),
+            close=Mock(),
+        ),
+    ):
         yield
 
 
@@ -174,24 +184,18 @@ def events(hass):
 @pytest.fixture
 async def mock_block_device():
     """Mock block (Gen1, CoAP) device."""
-    with patch("homeassistant.components.shelly.utils.COAP", autospec=True), patch(
+    with patch("aioshelly.block_device.BlockDevice.create") as block_device_mock:
-        "aioshelly.block_device.BlockDevice.create"
-    ) as block_device_mock:
 
         def update():
             block_device_mock.return_value.subscribe_updates.call_args[0][0]({})
 
         device = Mock(
+            spec=BlockDevice,
             blocks=MOCK_BLOCKS,
             settings=MOCK_SETTINGS,
             shelly=MOCK_SHELLY_COAP,
             status=MOCK_STATUS_COAP,
             firmware_version="some fw string",
-            update=AsyncMock(),
-            update_status=AsyncMock(),
-            trigger_ota_update=AsyncMock(),
-            trigger_reboot=AsyncMock(),
-            initialize=AsyncMock(),
             initialized=True,
         )
         block_device_mock.return_value = device
@@ -209,18 +213,13 @@ async def mock_rpc_device():
             rpc_device_mock.return_value.subscribe_updates.call_args[0][0]({})
 
         device = Mock(
-            call_rpc=AsyncMock(),
+            spec=RpcDevice,
             config=MOCK_CONFIG,
             event={},
             shelly=MOCK_SHELLY_RPC,
             status=MOCK_STATUS_RPC,
             firmware_version="some fw string",
-            update=AsyncMock(),
-            trigger_ota_update=AsyncMock(),
-            trigger_reboot=AsyncMock(),
-            initialize=AsyncMock(),
             initialized=True,
-            shutdown=AsyncMock(),
         )
 
         rpc_device_mock.return_value = device

tests/components/shelly/test_init.py

@@ -0,0 +1,184 @@
+"""Test cases for the Shelly component."""
+
+from unittest.mock import AsyncMock
+
+from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError
+import pytest
+
+from homeassistant.components.shelly.const import DOMAIN
+from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
+from homeassistant.const import STATE_ON, STATE_UNAVAILABLE
+from homeassistant.helpers import device_registry
+from homeassistant.setup import async_setup_component
+
+from . import MOCK_MAC, init_integration
+
+from tests.common import MockConfigEntry
+
+
+async def test_custom_coap_port(hass, mock_block_device, caplog):
+    """Test custom coap port."""
+    assert await async_setup_component(
+        hass,
+        DOMAIN,
+        {DOMAIN: {"coap_port": 7632}},
+    )
+    await hass.async_block_till_done()
+
+    await init_integration(hass, 1)
+    assert "Starting CoAP context with UDP port 7632" in caplog.text
+
+
+@pytest.mark.parametrize("gen", [1, 2])
+async def test_shared_device_mac(
+    hass, gen, mock_block_device, mock_rpc_device, device_reg, caplog
+):
+    """Test first time shared device with another domain."""
+    config_entry = MockConfigEntry(domain="test", data={}, unique_id="some_id")
+    config_entry.add_to_hass(hass)
+    device_reg.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={
+            (
+                device_registry.CONNECTION_NETWORK_MAC,
+                device_registry.format_mac(MOCK_MAC),
+            )
+        },
+    )
+    await init_integration(hass, gen, sleep_period=1000)
+    assert "will resume when device is online" in caplog.text
+
+
+async def test_setup_entry_not_shelly(hass, caplog):
+    """Test not Shelly entry."""
+    entry = MockConfigEntry(domain=DOMAIN, data={}, unique_id=DOMAIN)
+    entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(entry.entry_id) is False
+    await hass.async_block_till_done()
+
+    assert "probably comes from a custom integration" in caplog.text
+
+
+@pytest.mark.parametrize("gen", [1, 2])
+async def test_device_connection_error(
+    hass, gen, mock_block_device, mock_rpc_device, monkeypatch
+):
+    """Test device connection error."""
+    monkeypatch.setattr(
+        mock_block_device, "initialize", AsyncMock(side_effect=DeviceConnectionError)
+    )
+    monkeypatch.setattr(
+        mock_rpc_device, "initialize", AsyncMock(side_effect=DeviceConnectionError)
+    )
+
+    entry = await init_integration(hass, gen)
+    assert entry.state == ConfigEntryState.SETUP_RETRY
+
+
+@pytest.mark.parametrize("gen", [1, 2])
+async def test_device_auth_error(
+    hass, gen, mock_block_device, mock_rpc_device, monkeypatch
+):
+    """Test device authentication error."""
+    monkeypatch.setattr(
+        mock_block_device, "initialize", AsyncMock(side_effect=InvalidAuthError)
+    )
+    monkeypatch.setattr(
+        mock_rpc_device, "initialize", AsyncMock(side_effect=InvalidAuthError)
+    )
+
+    entry = await init_integration(hass, gen)
+    assert entry.state == ConfigEntryState.SETUP_ERROR
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+
+    flow = flows[0]
+    assert flow.get("step_id") == "reauth_confirm"
+    assert flow.get("handler") == DOMAIN
+
+    assert "context" in flow
+    assert flow["context"].get("source") == SOURCE_REAUTH
+    assert flow["context"].get("entry_id") == entry.entry_id
+
+
+@pytest.mark.parametrize("entry_sleep, device_sleep", [(None, 0), (1000, 1000)])
+async def test_sleeping_block_device_online(
+    hass, entry_sleep, device_sleep, mock_block_device, device_reg, caplog
+):
+    """Test sleeping block device online."""
+    config_entry = MockConfigEntry(domain=DOMAIN, data={}, unique_id="shelly")
+    config_entry.add_to_hass(hass)
+    device_reg.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections={
+            (
+                device_registry.CONNECTION_NETWORK_MAC,
+                device_registry.format_mac(MOCK_MAC),
+            )
+        },
+    )
+
+    entry = await init_integration(hass, 1, sleep_period=entry_sleep)
+    assert "will resume when device is online" in caplog.text
+
+    mock_block_device.mock_update()
+    assert "online, resuming setup" in caplog.text
+    assert entry.data["sleep_period"] == device_sleep
+
+
+@pytest.mark.parametrize("entry_sleep, device_sleep", [(None, 0), (1000, 1000)])
+async def test_sleeping_rpc_device_online(
+    hass, entry_sleep, device_sleep, mock_rpc_device, caplog
+):
+    """Test sleeping RPC device online."""
+    entry = await init_integration(hass, 2, sleep_period=entry_sleep)
+    assert "will resume when device is online" in caplog.text
+
+    mock_rpc_device.mock_update()
+    assert "online, resuming setup" in caplog.text
+    assert entry.data["sleep_period"] == device_sleep
+
+
+@pytest.mark.parametrize(
+    "gen, entity_id",
+    [
+        (1, "switch.test_name_channel_1"),
+        (2, "switch.test_switch_0"),
+    ],
+)
+async def test_entry_unload(hass, gen, entity_id, mock_block_device, mock_rpc_device):
+    """Test entry unload."""
+    entry = await init_integration(hass, gen)
+
+    assert entry.state is ConfigEntryState.LOADED
+    assert hass.states.get(entity_id).state is STATE_ON
+
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.NOT_LOADED
+    assert hass.states.get(entity_id).state is STATE_UNAVAILABLE
+
+
+@pytest.mark.parametrize(
+    "gen, entity_id",
+    [
+        (1, "switch.test_name_channel_1"),
+        (2, "switch.test_switch_0"),
+    ],
+)
+async def test_entry_unload_device_not_ready(
+    hass, gen, entity_id, mock_block_device, mock_rpc_device
+):
+    """Test entry unload when device is not ready."""
+    entry = await init_integration(hass, gen, sleep_period=1000)
+
+    assert entry.state is ConfigEntryState.LOADED
+    assert hass.states.get(entity_id) is None
+
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.NOT_LOADED

tests/components/shelly/test_update.py

@@ -6,7 +6,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_component import async_update_entity
 from homeassistant.helpers.entity_registry import async_get
 
-from . import init_integration
+from . import MOCK_MAC, init_integration
 
 
 async def test_block_update(hass: HomeAssistant, mock_block_device, monkeypatch):
@@ -15,7 +15,7 @@ async def test_block_update(hass: HomeAssistant, mock_block_device, monkeypatch)
     entity_registry.async_get_or_create(
         UPDATE_DOMAIN,
         DOMAIN,
-        "test-mac-fwupdate",
+        f"{MOCK_MAC}-fwupdate",
         suggested_object_id="test_name_firmware_update",
         disabled_by=None,
     )
@@ -46,7 +46,7 @@ async def test_rpc_update(hass: HomeAssistant, mock_rpc_device, monkeypatch):
     entity_registry.async_get_or_create(
         UPDATE_DOMAIN,
         DOMAIN,
-        "shelly-sys-fwupdate",
+        f"{MOCK_MAC}-sys-fwupdate",
         suggested_object_id="test_name_firmware_update",
         disabled_by=None,
     )