Replace asyncio.sleep with custom sleep function (#117)

Author: DeerMaximum

custom_components/ta_cmi/__init__.py

@@ -29,6 +29,21 @@
 PLATFORMS: list[str] = [Platform.SENSOR, Platform.BINARY_SENSOR]
 
 
+async def custom_sleep(delay: int) -> None:
+    """Custom sleep function to prevent Home Assistant from canceling."""
+    start = time.time()
+    try:
+        await asyncio.sleep(delay)
+    except asyncio.CancelledError:
+        elapsed = time.time() - start
+        _LOGGER.debug(
+            "Sleep cancelled after %s. Sleep remaining time: %s",
+            elapsed,
+            delay - elapsed,
+        )
+        await asyncio.sleep(delay - elapsed)
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up platform from a ConfigEntry."""
     host: str = entry.data[CONF_HOST]
@@ -106,21 +121,6 @@ def __init__(
 
         super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=update_interval)
 
-    @staticmethod
-    async def custom_sleep(delay: int) -> None:
-        """Custom sleep function to prevent Home Assistant from canceling."""
-        start = time.time()
-        try:
-            await asyncio.sleep(delay)
-        except asyncio.CancelledError:
-            elapsed = time.time() - start
-            _LOGGER.debug(
-                "Sleep cancelled after %s. Sleep remaining time: %s",
-                elapsed,
-                delay - elapsed,
-            )
-            await asyncio.sleep(delay - elapsed)
-
     async def _async_update_data(self) -> dict[str, Any]:
         """Update data."""
         try:
@@ -141,7 +141,7 @@ async def _async_update_data(self) -> dict[str, Any]:
                     _LOGGER.debug(
                         f"Sleep mode for {DEVICE_DELAY} seconds to prevent rate limiting"
                     )
-                    await self.custom_sleep(DEVICE_DELAY)
+                    await custom_sleep(DEVICE_DELAY)
 
             return return_data
         except (InvalidCredentialsError, RateLimitError, ApiError) as err:

custom_components/ta_cmi/config_flow.py

@@ -1,7 +1,6 @@
 """Config flow for Technische Alternative C.M.I. integration."""
 from __future__ import annotations
 
-import asyncio
 from copy import deepcopy
 from datetime import timedelta
 import time
@@ -15,16 +14,10 @@
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
-from ta_cmi import (
-    CMI,
-    ApiError,
-    Device,
-    InvalidCredentialsError,
-    InvalidDeviceError,
-    RateLimitError,
-)
+from ta_cmi import CMI, ApiError, Device, InvalidCredentialsError, RateLimitError
 import voluptuous as vol
 
+from . import custom_sleep
 from .const import (
     _LOGGER,
     CONF_CHANNELS,
@@ -69,13 +62,13 @@ async def fetch_device(device: Device, retry=False) -> None:
             _LOGGER.debug(
                 f"Sleep mode for {DEVICE_DELAY} seconds to prevent rate limiting"
             )
-            await asyncio.sleep(DEVICE_DELAY)
+            await custom_sleep(DEVICE_DELAY)
             device.set_device_type("DUMMY-NO-IO")
 
         _LOGGER.debug("Try to fetch device type: %s", device.id)
         await device.fetch_type()
         _LOGGER.debug(f"Sleep mode for {DEVICE_DELAY} seconds to prevent rate limiting")
-        await asyncio.sleep(DEVICE_DELAY)
+        await custom_sleep(DEVICE_DELAY)
 
         _LOGGER.debug("Try to fetch available device channels: %s", device.id)
         await device.update()
@@ -189,7 +182,7 @@ async def async_step_devices(
                 _LOGGER.debug(
                     f"Sleep mode for {DEVICE_DELAY} seconds to prevent rate limiting"
                 )
-                await asyncio.sleep(DEVICE_DELAY)
+                await custom_sleep(DEVICE_DELAY)
 
             try:
                 await fetch_device(dev)
@@ -291,7 +284,7 @@ async def async_step_finish(self) -> FlowResult:
         end_time = time.time()
         time_lapsed = end_time - self.start_time
         if time_lapsed <= DEVICE_DELAY:
-            await asyncio.sleep(DEVICE_DELAY - time_lapsed)
+            await custom_sleep(int(DEVICE_DELAY - time_lapsed))
         return self.async_create_entry(title="C.M.I", data=self.config)
 
     @staticmethod