fix: handle asyncio.CancelledError during tariff fetch to avoid setup… (#280)

* fix: handle asyncio.CancelledError during tariff fetch to avoid setup failure (iec-api aiohttp)\n\n- Catch CancelledError in _get_kwh_tariff and _get_kva_tariff\n- Fallback to 0.0 and warn instead of failing setup\n\nRefs: #278

* fix(config): map asyncio.CancelledError to cannot_connect in config flow; add broad handling around MFA and login\n\nfeat(coordinator): add fallback kWh/kVA tariffs from calculators API and handle CancelledError for customer/contracts fetch; return {} on cancelled refresh instead of failing\n\nAlso logs clearer messages; aligns with HA UX to avoid 'Unknown error occurred'.\n\nRefs: #278

---------

Co-authored-by: Guy Khmelnitsky <3136012+GuyKh@users.noreply.github.com>

Author: NirBY

custom_components/iec/config_flow.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import asyncio
 from collections.abc import Mapping
 from typing import Any
 
@@ -39,28 +40,34 @@ async def _validate_login(
     hass: HomeAssistant, login_data: dict[str, Any], api: IecClient
 ) -> dict[str, str]:
     """Validate login data and return any errors."""
-    assert login_data is not None
-    assert api is not None
-    assert login_data.get(CONF_USER_ID) is not None
-    assert (
-        login_data.get(CONF_TOTP_SECRET) or login_data.get(CONF_API_TOKEN) is not None
-    )
+    if not login_data or not api:
+        return {"base": "cannot_connect"}
+    if not login_data.get(CONF_USER_ID):
+        return {"base": "invalid_auth"}
+    if not (login_data.get(CONF_TOTP_SECRET) or login_data.get(CONF_API_TOKEN)):
+        return {"base": "invalid_auth"}
 
     if login_data.get(CONF_TOTP_SECRET):
         try:
             await api.verify_otp(login_data.get(CONF_TOTP_SECRET))
+        except asyncio.CancelledError:
+            return {"base": "cannot_connect"}
         except IECError:
             return {"base": "invalid_auth"}
 
     elif login_data.get(CONF_API_TOKEN):
         try:
             await api.load_jwt_token(JWT.from_dict(login_data.get(CONF_API_TOKEN)))
+        except asyncio.CancelledError:
+            return {"base": "cannot_connect"}
         except IECError:
             return {"base": "invalid_auth"}
 
     errors: dict[str, str] = {}
     try:
         await api.check_token()
+    except asyncio.CancelledError:
+        errors["base"] = "cannot_connect"
     except IECError:
         errors["base"] = "invalid_auth"
 
@@ -112,47 +119,78 @@ async def async_step_mfa(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
         """Handle MFA step."""
-        assert self.data is not None
-        assert self.data.get(CONF_USER_ID) is not None
+        if not self.data or not self.data.get(CONF_USER_ID):
+            return self.async_show_form(
+                step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors={"base": "invalid_auth"}
+            )
 
         client: IecClient = self.client
 
         errors: dict[str, str] = {}
         if user_input is not None and user_input.get(CONF_TOTP_SECRET) is not None:
-            data = {**self.data, **user_input}
-            errors = await _validate_login(self.hass, data, client)
-            if not errors:
-                data[CONF_API_TOKEN] = client.get_token().to_dict()
-
-                if data.get(CONF_TOTP_SECRET):
-                    data.pop(CONF_TOTP_SECRET)
-
-                customer = await client.get_customer()
-                data[CONF_BP_NUMBER] = customer.bp_number
-
-                contracts = await client.get_contracts(customer.bp_number)
-                contract_ids = [
-                    int(contract.contract_id)
-                    for contract in contracts
-                    if contract.status == 1
-                ]
-                if len(contract_ids) == 0:
-                    errors["base"] = "no_active_contracts"
-                elif len(contract_ids) == 1:
-                    data[CONF_SELECTED_CONTRACTS] = [contract_ids[0]]
-                    return self._async_create_iec_entry(data)
-                else:
-                    data[CONF_AVAILABLE_CONTRACTS] = contract_ids
-                    self.data = data
-                    return await self.async_step_select_contracts()
+            try:
+                data = {**self.data, **user_input}
+                errors = await _validate_login(self.hass, data, client)
+                if not errors:
+                    data[CONF_API_TOKEN] = client.get_token().to_dict()
+
+                    if data.get(CONF_TOTP_SECRET):
+                        data.pop(CONF_TOTP_SECRET)
+
+                    try:
+                        customer = await client.get_customer()
+                        data[CONF_BP_NUMBER] = customer.bp_number
+
+                        contracts = await client.get_contracts(customer.bp_number)
+                        contract_ids = [
+                            int(contract.contract_id)
+                            for contract in contracts
+                            if contract.status == 1
+                        ]
+                    except asyncio.CancelledError:
+                        errors["base"] = "cannot_connect"
+                    except IECError:
+                        errors["base"] = "cannot_connect"
+                    except Exception as err:  # noqa: BLE001
+                        _LOGGER.exception("Unexpected error during contracts fetch: %s", err)
+                        errors["base"] = "cannot_connect"
+
+                    if not errors:
+                        if len(contract_ids) == 0:
+                            errors["base"] = "no_active_contracts"
+                        elif len(contract_ids) == 1:
+                            data[CONF_SELECTED_CONTRACTS] = [contract_ids[0]]
+                            return self._async_create_iec_entry(data)
+                        else:
+                            data[CONF_AVAILABLE_CONTRACTS] = contract_ids
+                            self.data = data
+                            return await self.async_step_select_contracts()
+            except asyncio.CancelledError:
+                errors["base"] = "cannot_connect"
+            except IECError:
+                errors["base"] = "cannot_connect"
+            except Exception as err:  # noqa: BLE001
+                _LOGGER.exception("Unexpected error during MFA step: %s", err)
+                errors["base"] = "cannot_connect"
 
         if errors:
             schema = {vol.Required(CONF_USER_ID, default=self.data[CONF_USER_ID]): str}
         else:
             schema = {}
 
         schema[vol.Required(CONF_TOTP_SECRET)] = str
-        otp_type = await client.login_with_id()
+        try:
+            otp_type = await client.login_with_id()
+        except asyncio.CancelledError:
+            errors["base"] = errors.get("base") or "cannot_connect"
+            otp_type = "OTP"
+        except IECError:
+            errors["base"] = errors.get("base") or "cannot_connect"
+            otp_type = "OTP"
+        except Exception as err:  # noqa: BLE001
+            _LOGGER.exception("Unexpected error during login_with_id: %s", err)
+            errors["base"] = errors.get("base") or "cannot_connect"
+            otp_type = "OTP"
 
         return self.async_show_form(
             step_id="mfa",
@@ -243,7 +281,18 @@ async def async_step_reauth_confirm(
             )
             client = self.client
 
-        otp_type = await client.login_with_id()
+        try:
+            otp_type = await client.login_with_id()
+        except asyncio.CancelledError:
+            errors["base"] = errors.get("base") or "cannot_connect"
+            otp_type = "OTP"
+        except IECError:
+            errors["base"] = errors.get("base") or "cannot_connect"
+            otp_type = "OTP"
+        except Exception as err:  # noqa: BLE001
+            _LOGGER.exception("Unexpected error during reauth login_with_id: %s", err)
+            errors["base"] = errors.get("base") or "cannot_connect"
+            otp_type = "OTP"
 
         schema = {
             vol.Required(CONF_USER_ID): self.reauth_entry.data[CONF_USER_ID],

custom_components/iec/coordinator.py

@@ -1,6 +1,7 @@
 """Coordinator to handle IEC connections."""
 
 import calendar
+import asyncio
 import itertools
 import jwt
 import logging
@@ -189,20 +190,104 @@ async def _get_kwh_tariff(self) -> float:
         if not self._kwh_tariff:
             try:
                 self._kwh_tariff = await self.api.get_kwh_tariff()
+            except asyncio.CancelledError:
+                _LOGGER.warning(
+                    "Fetching kWh tariff was cancelled; using 0.0 and continuing"
+                )
+                self._kwh_tariff = 0.0
             except IECError as e:
                 _LOGGER.exception("Failed fetching kWh Tariff", e)
+            except Exception as e:
+                _LOGGER.exception("Unexpected error fetching kWh Tariff", e)
+
+            # Fallback: try IEC calculators API when main call failed or returned 0.0
+            if not self._kwh_tariff or self._kwh_tariff == 0.0:
+                kwh_fallback, _ = await self._fetch_tariffs_from_calculators()
+                if kwh_fallback and kwh_fallback > 0:
+                    _LOGGER.debug("Using fallback kWh tariff from calculators API: %s", kwh_fallback)
+                    self._kwh_tariff = kwh_fallback
         return self._kwh_tariff or 0.0
 
     async def _get_kva_tariff(self) -> float:
         if not self._kva_tariff:
             try:
                 self._kva_tariff = await self.api.get_kva_tariff()
+            except asyncio.CancelledError:
+                _LOGGER.warning(
+                    "Fetching kVA tariff was cancelled; using 0.0 and continuing"
+                )
+                self._kva_tariff = 0.0
             except IECError as e:
                 _LOGGER.exception("Failed fetching KVA Tariff from IEC API", e)
             except Exception as e:
-                _LOGGER.exception("Failed fetching KVA Tariff", e)
+                _LOGGER.exception("Unexpected error fetching KVA Tariff", e)
+
+            # Fallback: try IEC calculators API when main call failed or returned 0.0
+            if not self._kva_tariff or self._kva_tariff == 0.0:
+                _, kva_fallback = await self._fetch_tariffs_from_calculators()
+                if kva_fallback and kva_fallback > 0:
+                    _LOGGER.debug("Using fallback kVA tariff from calculators API: %s", kva_fallback)
+                    self._kva_tariff = kva_fallback
         return self._kva_tariff or 0.0
 
+    async def _fetch_tariffs_from_calculators(self) -> tuple[float | None, float | None]:
+        """Fetch tariffs from IEC calculators endpoints as a fallback.
+
+        Returns: tuple of (kwh_home_rate, kva_rate), each may be None if not found.
+        """
+        session = aiohttp_client.async_get_clientsession(self.hass, family=socket.AF_INET)
+        kwh_tariff: float | None = None
+        kva_tariff: float | None = None
+
+        # Primary fallback: calculators/period (contains both homeRate and kvaRate)
+        try:
+            async with session.get(
+                "https://iecapi.iec.co.il/api/content/he-IL/calculators/period",
+                timeout=30,
+            ) as resp:
+                if resp.status == 200:
+                    data = await resp.json(content_type=None)
+                    rates = data.get("period_Calculator_Rates") or {}
+                    kwh_val = rates.get("homeRate")
+                    kva_val = rates.get("kvaRate")
+                    if isinstance(kwh_val, (int, float)):
+                        kwh_tariff = float(kwh_val)
+                    if isinstance(kva_val, (int, float)):
+                        kva_tariff = float(kva_val)
+                    _LOGGER.debug(
+                        "Fetched fallback tariffs from calculators/period: homeRate=%s, kvaRate=%s",
+                        kwh_tariff,
+                        kva_tariff,
+                    )
+        except asyncio.CancelledError:
+            _LOGGER.debug("Fallback calculators/period fetch was cancelled")
+        except Exception as err:  # noqa: BLE001
+            _LOGGER.debug("Failed fetching fallback tariffs from calculators/period: %s", err)
+
+        # Secondary fallback: calculators/gadget (has homeRate only)
+        if kwh_tariff is None:
+            try:
+                async with session.get(
+                    "https://iecapi.iec.co.il/api/content/he-IL/calculators/gadget",
+                    timeout=30,
+                ) as resp:
+                    if resp.status == 200:
+                        data = await resp.json(content_type=None)
+                        rates = data.get("gadget_Calculator_Rates") or {}
+                        kwh_val = rates.get("homeRate")
+                        if isinstance(kwh_val, (int, float)):
+                            kwh_tariff = float(kwh_val)
+                        _LOGGER.debug(
+                            "Fetched fallback kWh tariff from calculators/gadget: homeRate=%s",
+                            kwh_tariff,
+                        )
+            except asyncio.CancelledError:
+                _LOGGER.debug("Fallback calculators/gadget fetch was cancelled")
+            except Exception as err:  # noqa: BLE001
+                _LOGGER.debug("Failed fetching fallback kWh tariff from calculators/gadget: %s", err)
+
+        return kwh_tariff, kva_tariff
+
     async def _get_delivery_tariff(self, phase) -> float:
         delivery_tariff = self._delivery_tariff_by_phase.get(phase)
         if not delivery_tariff:
@@ -373,10 +458,32 @@ async def _update_data(
         self,
     ) -> dict[str, dict[str, Any]]:
         if not self._bp_number:
-            customer = await self.api.get_customer()
-            self._bp_number = customer.bp_number
+            try:
+                customer = await self.api.get_customer()
+                self._bp_number = customer.bp_number
+            except asyncio.CancelledError:
+                _LOGGER.warning(
+                    "Fetching customer was cancelled; using empty BP number and skipping contracts"
+                )
+                self._bp_number = None
+            except IECError as e:
+                _LOGGER.exception("Failed fetching customer", e)
+                self._bp_number = None
 
-        all_contracts: list[Contract] = await self.api.get_contracts(self._bp_number)
+        try:
+            all_contracts: list[Contract] = (
+                await self.api.get_contracts(self._bp_number)
+                if self._bp_number
+                else []
+            )
+        except asyncio.CancelledError:
+            _LOGGER.warning(
+                "Fetching contracts was cancelled; continuing with empty contracts"
+            )
+            all_contracts = []
+        except IECError as e:
+            _LOGGER.exception("Failed fetching contracts", e)
+            all_contracts = []
         if not self._contract_ids:
             self._contract_ids = [
                 int(contract.contract_id)
@@ -428,6 +535,11 @@ async def _update_data(
                 billing_invoices = await self.api.get_billing_invoices(
                     self._bp_number, contract_id
                 )
+            except asyncio.CancelledError:
+                _LOGGER.warning(
+                    "Fetching invoices was cancelled; continuing without invoices"
+                )
+                billing_invoices = None
             except IECError as e:
                 _LOGGER.exception("Failed fetching invoices", e)
                 billing_invoices = None
@@ -669,6 +781,13 @@ async def _async_update_data(
 
         try:
             return await self._update_data()
+        except asyncio.CancelledError as err:
+            _LOGGER.warning(
+                "Data update was cancelled (network timeout/cancelled); will retry later: %s",
+                err,
+            )
+            # Return empty data so setup doesn't fail; periodic refresh will try again
+            return {}
         except Exception as err:
             _LOGGER.error("Failed updating data. Exception: %s", err)
             _LOGGER.error(traceback.format_exc())