Raise config not ready exception.

custom_components/plum_ecomax/__init__.py

@@ -7,6 +7,7 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryNotReady
 from pyplumio.helpers.filters import delta
 from pyplumio.structures.alerts import Alert
 
@@ -53,7 +54,10 @@ async def async_close_connection(event=None):
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_close_connection)
     )
 
-    load_ok = await connection.async_setup()
+    load_ok, error_message = await connection.async_setup()
+    if not load_ok:
+        raise ConfigEntryNotReady(error_message)
+
     await async_setup_services(hass, connection)
     await async_setup_events(hass, connection)
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = connection
@@ -84,7 +88,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
         try:
-            connection = hass.data[DOMAIN][entry.entry_id]
+            connection: EcomaxConnection = hass.data[DOMAIN][entry.entry_id]
             await connection.close()
             hass.data[DOMAIN].pop(entry.entry_id)
         except KeyError:

custom_components/plum_ecomax/connection.py

@@ -4,7 +4,7 @@
 import asyncio
 from collections.abc import Mapping
 import logging
-from typing import Any, Final
+from typing import Any, Final, Tuple
 
 from homeassistant.components.network import async_get_source_ip
 from homeassistant.components.network.const import IPV4_BROADCAST_ADDR
@@ -116,19 +116,21 @@ def __getattr__(self, name: str):
 
         raise AttributeError()
 
-    async def async_setup(self) -> bool:
+    async def async_setup(self) -> Tuple[bool, str | None]:
         """Setup connection and add hass stop handler."""
 
+        error_message = None
         await self.connection.connect()
         try:
             self._device = await self.connection.get_device(
-                "ecomax", timeout=DEVICE_TIMEOUT
+                ECOMAX, timeout=DEVICE_TIMEOUT
             )
         except asyncio.TimeoutError:
-            _LOGGER.error("Device has failed to respond in time")
-            return False
+            error_message = "Device has failed to respond in time"
+            await self.connection.close()
+            return False, error_message
 
-        return True
+        return True, error_message
 
     async def async_update_device_capabilities(self) -> None:
         """Update device capabilities."""

tests/test_connection.py

@@ -114,7 +114,7 @@ async def test_async_setup(
     assert await connection.async_setup()
 
     mock_connection_handler.connect.assert_awaited_once()
-    mock_connection_handler.get_device.assert_awaited_once_with(ECOMAX, timeout=10)
+    mock_connection_handler.get_device.assert_awaited_once_with(ECOMAX, timeout=20)
 
     # Check connection class properties.
     assert connection.host == "localhost"
@@ -135,8 +135,9 @@ async def test_async_setup(
     )
 
     # Check with device timeout.
-    assert not await connection.async_setup()
-    assert "Device has failed to respond in time" in caplog.text
+    load_ok, error_message = await connection.async_setup()
+    assert not load_ok
+    assert error_message == "Device has failed to respond in time"
 
     # Check name with serial connection.
     mock_connection_handler = AsyncMock(spec=SerialConnection)

tests/test_init.py

@@ -6,8 +6,10 @@
 from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.util import dt as dt_util
 from pyplumio.structures.alerts import Alert
+import pytest
 
 from custom_components.plum_ecomax import (
     async_migrate_entry,
@@ -26,7 +28,10 @@
 )
 
 
-@patch("custom_components.plum_ecomax.EcomaxConnection.async_setup")
+@patch(
+    "custom_components.plum_ecomax.EcomaxConnection.async_setup",
+    side_effect=((True, None), (False, "error message")),
+)
 @patch("custom_components.plum_ecomax.async_setup_services")
 @patch.object(EcomaxConnection, "close", create=True, new_callable=AsyncMock)
 async def test_setup_and_unload_entry(
@@ -41,10 +46,14 @@ async def test_setup_and_unload_entry(
     assert DOMAIN in hass.data and config_entry.entry_id in hass.data[DOMAIN]
     assert isinstance(hass.data[DOMAIN][config_entry.entry_id], EcomaxConnection)
 
+    # Test with exception.
+    with pytest.raises(ConfigEntryNotReady):
+        await async_setup_entry(hass, config_entry)
+
     # Send HA stop event and check that connection was closed.
     hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
     await hass.async_block_till_done()
-    mock_close.assert_awaited_once()
+    assert mock_close.call_count == 2
     mock_close.reset_mock()
 
     # Unload entry and verify that it is no longer present in hass data.