Rework reconnection logic, add better connection error handling

Author: GnoX

custom_components/ef_ble/__init__.py

@@ -1,19 +1,36 @@
 """The unofficial EcoFlow BLE devices integration"""
 
-from __future__ import annotations
-
 import logging
+from functools import partial
 
+import homeassistant.helpers.issue_registry as ir
 from homeassistant.components import bluetooth
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ADDRESS, CONF_TYPE, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.exceptions import (
+    ConfigEntryError,
+    ConfigEntryNotReady,
+)
 from homeassistant.helpers.device_registry import DeviceInfo
 
 from . import eflib
 from .config_flow import ConfLogOptions, LogOptions
-from .const import CONF_UPDATE_PERIOD, CONF_USER_ID, DOMAIN, MANUFACTURER
+from .const import (
+    CONF_CONNECTION_TIMEOUT,
+    CONF_UPDATE_PERIOD,
+    CONF_USER_ID,
+    DEFAULT_CONNECTION_TIMEOUT,
+    DEFAULT_UPDATE_PERIOD,
+    DOMAIN,
+    MANUFACTURER,
+)
+from .eflib.connection import (
+    AuthFailedError,
+    BleakError,
+    ConnectionTimeout,
+    MaxConnectionAttemptsReached,
+)
 
 PLATFORMS: list[Platform] = [
     Platform.SENSOR,
@@ -27,18 +44,8 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-
-class _ConfigNotReady(ConfigEntryNotReady):
-    def __init__(
-        self,
-        translation_key: str | None = None,
-        translation_placeholders: dict[str, str] | None = None,
-    ) -> None:
-        super().__init__(
-            translation_domain=DOMAIN,
-            translation_key=translation_key,
-            translation_placeholders=translation_placeholders,
-        )
+ConfigEntryNotReady = partial(ConfigEntryNotReady, translation_domain=DOMAIN)
+ConfigEntryError = partial(ConfigEntryError, translation_domain=DOMAIN)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: DeviceConfigEntry) -> bool:
@@ -49,45 +56,91 @@ async def async_setup_entry(hass: HomeAssistant, entry: DeviceConfigEntry) -> bo
     address = entry.data.get(CONF_ADDRESS)
     user_id = entry.data.get(CONF_USER_ID)
     merged_options = entry.data | entry.options
-    update_period = merged_options.get(CONF_UPDATE_PERIOD, 0)
+    update_period = merged_options.get(CONF_UPDATE_PERIOD, DEFAULT_UPDATE_PERIOD)
+    timeout = merged_options.get(CONF_CONNECTION_TIMEOUT, DEFAULT_CONNECTION_TIMEOUT)
 
     if address is None or user_id is None:
         return False
 
     if not bluetooth.async_address_present(hass, address):
-        raise _ConfigNotReady("device_not_present")
+        raise ConfigEntryNotReady(translation_key="device_not_present")
 
     _LOGGER.debug("Connecting Device")
-    discovery_info = bluetooth.async_last_service_info(hass, address, connectable=True)
-    device = eflib.NewDevice(discovery_info.device, discovery_info.advertisement)
+    device: eflib.DeviceBase | None = getattr(entry, "runtime_data", None)
     if device is None:
-        raise _ConfigNotReady("unable_to_create_device")
+        discovery_info = bluetooth.async_last_service_info(
+            hass, address, connectable=True
+        )
+        device = eflib.NewDevice(discovery_info.device, discovery_info.advertisement)
+        if device is None:
+            raise ConfigEntryNotReady(translation_key="unable_to_create_device")
 
-    await (
-        device.with_update_period(update_period)
-        .with_logging_options(ConfLogOptions.from_config(merged_options))
-        .connect(user_id)
-    )
-    entry.runtime_data = device
+        entry.runtime_data = device
 
-    timeout = 30
-    state = await device.wait_until_connected_or_error(timeout=timeout)
+    issue_id = f"{entry.entry_id}_max_connection_attempts"
 
-    if state.connection_error():
-        raise _ConfigNotReady(
-            "could_not_connect", translation_placeholders={"time": str(timeout)}
+    try:
+        await (
+            device.with_update_period(update_period)
+            .with_logging_options(ConfLogOptions.from_config(merged_options))
+            .with_disabled_reconnect()
+            .connect(user_id, timeout=timeout)
         )
-    if state.is_error():
-        raise _ConfigNotReady("error_after_connected")
-    if not state.authenticated():
-        raise _ConfigNotReady("could_not_authenticate")
+        state = await device.wait_until_authenticated_or_error(raise_on_error=True)
+    except (ConnectionTimeout, BleakError, TimeoutError) as e:
+        raise ConfigEntryNotReady(
+            translation_key="could_not_connect",
+            translation_placeholders={"time": str(timeout)},
+        ) from e
+    except AuthFailedError as e:
+        raise ConfigEntryNotReady(translation_key="authentication_failed") from e
+    except MaxConnectionAttemptsReached as e:
+        await device.disconnect()
+        ir.async_create_issue(
+            hass,
+            DOMAIN,
+            issue_id,
+            is_fixable=False,
+            severity=ir.IssueSeverity.ERROR,
+            translation_key="max_connection_attempts_reached",
+            translation_placeholders={
+                "device_name": device.name,
+                "attempts": str(e.attempts),
+            },
+        )
+        raise ConfigEntryError(
+            translation_key="could_not_connect_no_retry",
+            translation_placeholders={"attempts": str(e.attempts)},
+        ) from e
+    except Exception as e:
+        _LOGGER.exception("Unknown error")
+        await device.disconnect()
+        raise ConfigEntryNotReady(
+            translation_key="unknown_error", translation_placeholders={"error": str(e)}
+        ) from e
+    else:
+        if not state.authenticated():
+            await device.disconnect()
+            raise ConfigEntryNotReady(
+                translation_key="failed_after_successful_connection",
+                translation_placeholders={"last_state": state},
+            )
+    ir.async_delete_issue(hass, DOMAIN, issue_id)
 
     _LOGGER.debug("Creating entities")
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     _LOGGER.debug("Setup done")
     entry.async_on_unload(entry.add_update_listener(_update_listener))
 
+    def _on_disconnect(exc: Exception | type[Exception] | None):
+        async def _disconnect_and_reload():
+            hass.config_entries.async_schedule_reload(entry.entry_id)
+
+        hass.async_create_task(_disconnect_and_reload())
+
+    entry.async_on_unload(device.on_disconnect(_on_disconnect))
+
     return True
 
 
@@ -112,7 +165,7 @@ def device_info(entry: ConfigEntry) -> DeviceInfo:
 async def _update_listener(hass: HomeAssistant, entry: DeviceConfigEntry):
     device = entry.runtime_data
     merged_options = entry.data | entry.options
-    update_period = merged_options.get(CONF_UPDATE_PERIOD, 0)
-    device.with_update_period(update_period).with_logging_options(
+    update_period = merged_options.get(CONF_UPDATE_PERIOD, DEFAULT_UPDATE_PERIOD)
+    device.with_update_period(period=update_period).with_logging_options(
         ConfLogOptions.from_config(merged_options)
     )

custom_components/ef_ble/config_flow.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import base64
 import logging
 from collections.abc import Mapping
@@ -28,6 +29,7 @@
 
 from . import eflib
 from .const import (
+    CONF_CONNECTION_TIMEOUT,
     CONF_LOG_BLEAK,
     CONF_LOG_CONNECTION,
     CONF_LOG_ENCRYPTED_PAYLOADS,
@@ -37,6 +39,8 @@
     CONF_LOG_PAYLOADS,
     CONF_UPDATE_PERIOD,
     CONF_USER_ID,
+    DEFAULT_CONNECTION_TIMEOUT,
+    DEFAULT_UPDATE_PERIOD,
     DOMAIN,
 )
 from .eflib.connection import ConnectionState
@@ -117,6 +121,7 @@ async def async_step_bluetooth_confirm(
                     **self._login_option(),
                     vol.Required(CONF_ADDRESS): vol.In([f"{title} ({device.address})"]),
                     **_update_period_option(),
+                    **_timeout_option(),
                     **ConfLogOptions.schema(
                         ConfLogOptions.to_config(self._log_options)
                     ),
@@ -176,6 +181,7 @@ async def async_step_user(
                     **self._login_option(),
                     vol.Required(CONF_ADDRESS): vol.In(self._discovered_devices.keys()),
                     **_update_period_option(),
+                    **_timeout_option(),
                     **ConfLogOptions.schema(
                         ConfLogOptions.to_config(self._log_options)
                     ),
@@ -242,6 +248,7 @@ async def _validate_user_id(
         self._email = user_input.get("login", {}).get(CONF_EMAIL, "")
         password = user_input.get("login", {}).get(CONF_PASSWORD, "")
         user_id = user_input.get(CONF_USER_ID, "")
+        timeout = user_input.get(CONF_CONNECTION_TIMEOUT, 20)
 
         self._collapsed = False
 
@@ -259,8 +266,10 @@ async def _validate_user_id(
 
         device.with_logging_options(ConfLogOptions.from_config(user_input))
 
-        await device.connect(self._user_id, max_attempts=4)
-        conn_state = await device.wait_until_connected_or_error(timeout=20)
+        await device.connect(self._user_id)
+        conn_state = await asyncio.wait_for(
+            device.wait_until_authenticated_or_error(), timeout
+        )
         await device.disconnect()
 
         error = None
@@ -335,7 +344,9 @@ async def async_step_init(self, user_input: dict[str, Any] | None = None):
 
         merged_entry = self.config_entry.data | self.config_entry.options
         options = {
-            CONF_UPDATE_PERIOD: merged_entry.get(CONF_UPDATE_PERIOD, 0),
+            CONF_UPDATE_PERIOD: merged_entry.get(
+                CONF_UPDATE_PERIOD, DEFAULT_UPDATE_PERIOD
+            ),
         }
 
         device: eflib.DeviceBase | None = getattr(
@@ -412,9 +423,17 @@ def schema(
         }
 
 
-def _update_period_option(default: int = 0):
+def _update_period_option(default: int = DEFAULT_UPDATE_PERIOD):
     return {
         vol.Optional(CONF_UPDATE_PERIOD, default=default): vol.All(
             int, vol.Range(min=0)
         )
     }
+
+
+def _timeout_option(default: int = DEFAULT_CONNECTION_TIMEOUT):
+    return {
+        vol.Optional(CONF_CONNECTION_TIMEOUT, default=default): vol.All(
+            int, vol.Range(min=0)
+        )
+    }

custom_components/ef_ble/const.py

@@ -5,6 +5,7 @@
 
 CONF_USER_ID = "user_id"
 CONF_UPDATE_PERIOD = "update_period"
+CONF_CONNECTION_TIMEOUT = "connection_timeout"
 
 CONF_LOG_MASKED = "log_masked"
 CONF_LOG_PACKETS = "log_packets"
@@ -13,3 +14,7 @@
 CONF_LOG_MESSAGES = "log_messages"
 CONF_LOG_CONNECTION = "log_connection"
 CONF_LOG_BLEAK = "log_bleak"
+
+
+DEFAULT_UPDATE_PERIOD = 10
+DEFAULT_CONNECTION_TIMEOUT = 20