refactor: another refactor

Signed-off-by: Rafal Zielinski <sq4ind@gmail.com>

custom_components/adguard_hub/__init__.py

@@ -6,7 +6,7 @@ Transform your AdGuard Home into a smart network management powerhouse.
 import asyncio
 import logging
 from datetime import timedelta
-from typing import Dict, Any
+from typing import Any, Dict
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
@@ -15,8 +15,8 @@ from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .api import AdGuardHomeAPI, AdGuardConnectionError
-from .const import DOMAIN, PLATFORMS, SCAN_INTERVAL, CONF_SSL, CONF_VERIFY_SSL
+from .api import AdGuardHomeAPI, AdGuardConnectionError, AdGuardHomeError
+from .const import CONF_SSL, CONF_VERIFY_SSL, DOMAIN, PLATFORMS, SCAN_INTERVAL
 from .services import AdGuardControlHubServices
 
 _LOGGER = logging.getLogger(__name__)
@@ -43,19 +43,22 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
         _LOGGER.info(
             "Successfully connected to AdGuard Home at %s:%s",
-            entry.data[CONF_HOST], 
+            entry.data[CONF_HOST],
             entry.data[CONF_PORT]
         )
-    except Exception as err:
+    except AdGuardHomeError as err:
         _LOGGER.error("Failed to connect to AdGuard Home: %s", err)
         raise ConfigEntryNotReady(f"Unable to connect: {err}") from err
+    except Exception as err:
+        _LOGGER.exception("Unexpected error connecting to AdGuard Home")
+        raise ConfigEntryNotReady(f"Unexpected error: {err}") from err
 
     # Create update coordinator
     coordinator = AdGuardControlHubCoordinator(hass, api)
 
     try:
         await coordinator.async_config_entry_first_refresh()
-    except Exception as err:
+    except UpdateFailed as err:
         _LOGGER.error("Failed to perform initial data refresh: %s", err)
         raise ConfigEntryNotReady(f"Failed to fetch initial data: {err}") from err
 
@@ -72,17 +75,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     except Exception as err:
         _LOGGER.error("Failed to set up platforms: %s", err)
         # Clean up on failure
-        hass.data[DOMAIN].pop(entry.entry_id)
+        hass.data[DOMAIN].pop(entry.entry_id, None)
         raise ConfigEntryNotReady(f"Failed to set up platforms: {err}") from err
 
     # Register services (only once)
-    if not hass.services.has_service(DOMAIN, "block_services"):
+    services_key = f"{DOMAIN}_services"
+    if services_key not in hass.data:
         services = AdGuardControlHubServices(hass)
         services.register_services()
-        hass.data.setdefault(f"{DOMAIN}_services", services)
+        hass.data[services_key] = services
 
-    _LOGGER.info("AdGuard Control Hub setup complete for %s:%s", 
-                 entry.data[CONF_HOST], entry.data[CONF_PORT])
+    _LOGGER.info("AdGuard Control Hub setup complete")
     return True
 
 
@@ -92,14 +95,15 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     if unload_ok:
         # Remove this entry's data
-        hass.data[DOMAIN].pop(entry.entry_id)
+        hass.data[DOMAIN].pop(entry.entry_id, None)
 
         # Unregister services if this was the last entry
         if not hass.data[DOMAIN]:
-            services = hass.data.get(f"{DOMAIN}_services")
+            services_key = f"{DOMAIN}_services"
+            services = hass.data.get(services_key)
             if services:
                 services.unregister_services()
-                hass.data.pop(f"{DOMAIN}_services", None)
+                hass.data.pop(services_key, None)
             hass.data.pop(DOMAIN, None)
 
     return unload_ok
@@ -108,7 +112,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 class AdGuardControlHubCoordinator(DataUpdateCoordinator):
     """AdGuard Control Hub data update coordinator."""
 
-    def __init__(self, hass: HomeAssistant, api: AdGuardHomeAPI):
+    def __init__(self, hass: HomeAssistant, api: AdGuardHomeAPI) -> None:
         """Initialize the coordinator."""
         super().__init__(
             hass,
@@ -134,7 +138,7 @@ class AdGuardControlHubCoordinator(DataUpdateCoordinator):
             results = await asyncio.gather(*tasks, return_exceptions=True)
             clients, statistics, status = results
 
-            # Update stored data (use empty dict if fetch failed)
+            # Update stored data (use previous data if fetch failed)
             if not isinstance(clients, Exception):
                 self._clients = {
                     client["name"]: client