Refactor connect/disconnect (#247)

* Refactor init code

* Add init tests

Author: RobBie1221

custom_components/nefiteasy/__init__.py

@@ -43,14 +43,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     client = NefitEasy(hass, credentials)
 
     await client.connect()
-    _LOGGER.debug("Is connected state? %s", client.connected_state)
-
     if client.connected_state == STATE_CONNECTION_VERIFIED:
         hass.data[DOMAIN][entry.entry_id]["client"] = client
-        _LOGGER.info(
-            "Successfully connected %s to Nefit device!",
-            credentials.get(CONF_SERIAL),
-        )
     else:
         raise ConfigEntryNotReady
 
@@ -66,23 +60,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload nefit easy component."""
-    if not all(
-        await asyncio.gather(
-            *(
-                hass.config_entries.async_forward_entry_unload(entry, component)
-                for component in DOMAINS
-            )
-        )
-    ):
-        return False
+    unload_ok = await hass.config_entries.async_unload_platforms(entry, DOMAINS)
 
-    client = hass.data[DOMAIN][entry.entry_id]["client"]
+    if unload_ok:
+        client = hass.data[DOMAIN][entry.entry_id]["client"]
 
-    await client.shutdown("Unload entry")
+        await client.shutdown("Unload entry")
 
-    hass.data[DOMAIN].pop(entry.entry_id)
+        hass.data[DOMAIN].pop(entry.entry_id)
 
-    return True
+    return unload_ok
 
 
 class NefitEasy(DataUpdateCoordinator):
@@ -139,85 +126,49 @@ async def add_key(self, entity_description: NefitEntityDescription) -> None:
 
     async def connect(self) -> None:
         """Connect to nefit easy."""
-        _LOGGER.debug("Starting connecting..")
+        _LOGGER.debug("Start connecting.")
         if not self.is_connecting:
             self.is_connecting = True
-            retries_connection = 0
 
-            while self.connected_state != STATE_CONNECTED and retries_connection < 3:
-                await self.nefit.connect()
-                _LOGGER.debug("Waiting for connected event")
+            await self.nefit.connect()
+            _LOGGER.debug("Waiting for connected event.")
+            try:
+                await asyncio.wait_for(
+                    self.nefit.xmppclient.connected_event.wait(), timeout=29.0
+                )
+            except asyncio.TimeoutError:
+                _LOGGER.debug("TimeoutError on waiting for connected event.")
+            except:  # noqa: E722 pylint: disable=bare-except
+                _LOGGER.debug("Unknown error.")
+            else:
+                _LOGGER.debug("Connected successfully.")
+                self.connected_state = STATE_CONNECTED
+
+            if self.connected_state == STATE_CONNECTED:
+                self.nefit.get("/gateway/brandID")
                 try:
                     await asyncio.wait_for(
-                        self.nefit.xmppclient.connected_event.wait(), timeout=29.0
-                    )
-                    self.connected_state = STATE_CONNECTED
-                    _LOGGER.debug("adding stop listener")
-                    self.hass.bus.async_listen_once(
-                        EVENT_HOMEASSISTANT_STOP, self.shutdown
+                        self.nefit.xmppclient.message_event.wait(), timeout=29.0
                     )
                 except asyncio.TimeoutError:
                     _LOGGER.debug(
-                        "TimeoutError on waiting for connected event (connection retries=%d)",
-                        retries_connection,
+                        "Did not get a response in time for testing connection."
                     )
-                    retries_connection = retries_connection + 1
                 except:  # noqa: E722 pylint: disable=bare-except
-                    _LOGGER.debug("Unknown error")
-
-            # test password for decrypting messages if connected
-            if self.connected_state == STATE_CONNECTED:
-                _LOGGER.info(
-                    "Testing connection (connect retries=%d)", retries_connection
-                )
-                retries_validation = 0
-                while (
-                    self.connected_state != STATE_CONNECTION_VERIFIED
-                    and retries_validation < 3
-                ):
-                    try:
-                        self.nefit.get("/gateway/brandID")
-                        await asyncio.wait_for(
-                            self.nefit.xmppclient.message_event.wait(), timeout=29.0
-                        )
-                        self.nefit.xmppclient.message_event.clear()
-
-                        if self.connected_state == STATE_ERROR_AUTH:
-                            self.is_connecting = False
-                            return
+                    _LOGGER.debug("No connection while testing connection.")
+                else:
+                    self.nefit.xmppclient.message_event.clear()
 
+                    # No exception and no auth error
+                    if self.connected_state == STATE_CONNECTED:
                         self.connected_state = STATE_CONNECTION_VERIFIED
-                        _LOGGER.info(
-                            "Connected %s with %d retries and %d test retries.",
-                            self.serial,
-                            retries_connection,
-                            retries_validation,
-                        )
-                    except asyncio.TimeoutError:
-                        _LOGGER.error(
-                            "Did not get a response in time for testing connection (validation retries=%d).",
-                            retries_validation,
-                        )
-                        retries_validation = retries_validation + 1
-                    except:  # noqa: E722 pylint: disable=bare-except
-                        _LOGGER.error("No connection while testing connection.")
-                        break
 
             if self.connected_state != STATE_CONNECTION_VERIFIED:
-                self.hass.components.persistent_notification.create(
-                    f"Did not succeed in connecting {self.serial} to Bosch cloud after retrying 3 times. Retry in 30 seconds.",
-                    title="Nefit connect error",
-                    notification_id="nefit_connect_error",
-                )
-                self.is_connecting = False
-
-                # wait 30 seconds to retry
-                await asyncio.sleep(30)
-                await self.connect()
+                _LOGGER.debug("Successfully verified connection.")
 
             self.is_connecting = False
         else:
-            _LOGGER.debug("Connection procedure was already running..")
+            _LOGGER.debug("Connection procedure is already running.")
 
     async def shutdown(self, event: str) -> None:
         """Shutdown."""
@@ -250,23 +201,20 @@ async def failed_auth_handler(self, event: str) -> None:
                 notification_id="nefit_logon_error",
             )
 
+        if self.config_entry:
+            self.config_entry.async_start_reauth(self.hass)
+
     async def session_end_callback(self) -> None:
         """If connection is closed unexpectedly, try to reconnect."""
         if not self.expected_end:
-            self.hass.components.persistent_notification.create(
-                f"Unexpected disconnect of {self.serial} with Bosch server. Try to reconnect..",
-                title="Nefit disconnect",
-                notification_id="nefit_disconnect",
+            _LOGGER.debug(
+                f"Unexpected disconnect of {self.serial} with Bosch server. Try to reconnect.."
             )
 
-            _LOGGER.info("Starting reconnect procedure.")
             # Reset values
             self.connected_state = STATE_INIT
             self.expected_end = False
 
-            # Retry connection
-            await self.connect()
-
     async def parse_message(self, data: dict[str, Any]) -> None:
         """Message received callback function for the XMPP client."""
         if (
@@ -310,7 +258,10 @@ async def parse_message(self, data: dict[str, Any]) -> None:
     async def _async_update_data(self) -> dict[str, Any]:
         """Update data via library."""
         if self.connected_state != STATE_CONNECTION_VERIFIED:
-            raise UpdateFailed("Nefit easy not connected!")
+            _LOGGER.debug("Starting reconnect procedure.")
+            await self.connect()
+            if self.connected_state != STATE_CONNECTION_VERIFIED:
+                raise UpdateFailed("Nefit easy not connected!")
 
         async with self._lock:
             _url = "/ecus/rrc/uiStatus"

tests/test_init.py

@@ -1,4 +1,5 @@
 """Tests of the initialization of the nefiteasy integration."""
+import asyncio
 from unittest.mock import patch
 
 from homeassistant import config_entries
@@ -7,30 +8,75 @@
 
 from .conftest import ClientMock
 
+entry_data = {
+    "serial": "123456789",
+    "accesskey": "myAccessKey",
+    "password": "myPass",
+    "min_temp": 10,
+    "max_temp": 28,
+    "temp_step": 0.5,
+    "name": "Nefit",
+}
+
 
 @patch("custom_components.nefiteasy.NefitCore")
-async def test_setup_entry(mock_class, hass: HomeAssistant):
+async def test_load_unload_entry(mock_class, hass: HomeAssistant):
     """Validate that setup entry also configure the client."""
     client = ClientMock(mock_class)
     mock_class.return_value = client
 
-    config_entry = MockConfigEntry(
-        domain="nefiteasy",
-        data={
-            "serial": "123456789",
-            "accesskey": "myAccessKey",
-            "password": "myPass",
-            "min_temp": 10,
-            "max_temp": 28,
-            "temp_step": 0.5,
-            "name": "Nefit",
-        },
-    )
+    config_entry = MockConfigEntry(domain="nefiteasy", data=entry_data)
 
     config_entry.add_to_hass(hass)
 
     assert await hass.config_entries.async_setup(config_entry.entry_id)
-
     await hass.async_block_till_done()
 
     assert config_entry.state == config_entries.ConfigEntryState.LOADED
+
+    assert await hass.config_entries.async_unload(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert config_entry.state == config_entries.ConfigEntryState.NOT_LOADED
+
+
+@patch("custom_components.nefiteasy.NefitCore")
+async def test_setup_connection_fail_timeout(mock_class, hass: HomeAssistant):
+    """Test setup connection with timeout failure."""
+    client = ClientMock(mock_class)
+    mock_class.return_value = client
+
+    config_entry = MockConfigEntry(domain="nefiteasy", data=entry_data)
+
+    config_entry.add_to_hass(hass)
+
+    async def wait():
+        raise asyncio.TimeoutError
+
+    client.xmppclient.connected_event.wait = wait
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert config_entry.state == config_entries.ConfigEntryState.SETUP_RETRY
+
+
+@patch("custom_components.nefiteasy.NefitCore")
+async def test_setup_validation_fail_timeout(mock_class, hass: HomeAssistant):
+    """Test setup with validation timeout."""
+    client = ClientMock(mock_class)
+    mock_class.return_value = client
+
+    config_entry = MockConfigEntry(domain="nefiteasy", data=entry_data)
+
+    config_entry.add_to_hass(hass)
+
+    async def wait():
+        raise asyncio.TimeoutError
+
+    client.xmppclient.message_event.wait = wait
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert config_entry.state == config_entries.ConfigEntryState.SETUP_RETRY