fix: multiple fixes

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

custom_components/adguard_hub/config_flow.py

@@ -1,73 +1,128 @@
 """Config flow for AdGuard Control Hub integration."""
+import asyncio
 import logging
-from typing import Any
+from typing import Any, Dict, Optional
+
 import voluptuous as vol
 from homeassistant import config_entries
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from .api import AdGuardHomeAPI
-from .const import CONF_SSL, CONF_VERIFY_SSL, DEFAULT_PORT, DEFAULT_SSL, DEFAULT_VERIFY_SSL, DOMAIN
+from homeassistant.data_entry_flow import FlowResult
+import homeassistant.helpers.config_validation as cv
+
+from .api import AdGuardHomeAPI, AdGuardConnectionError, AdGuardAuthError
+from .const import (
+    CONF_SSL,
+    CONF_VERIFY_SSL,
+    DEFAULT_PORT,
+    DEFAULT_SSL,
+    DEFAULT_VERIFY_SSL,
+    DOMAIN,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
 STEP_USER_DATA_SCHEMA = vol.Schema({
-    vol.Required(CONF_HOST): str,
-    vol.Optional(CONF_PORT, default=DEFAULT_PORT): int,
-    vol.Optional(CONF_USERNAME): str,
-    vol.Optional(CONF_PASSWORD): str,
-    vol.Optional(CONF_SSL, default=DEFAULT_SSL): bool,
-    vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): bool,
+    vol.Required(CONF_HOST): cv.string,
+    vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
+    vol.Optional(CONF_USERNAME): cv.string,
+    vol.Optional(CONF_PASSWORD): cv.string,
+    vol.Optional(CONF_SSL, default=DEFAULT_SSL): cv.boolean,
+    vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): cv.boolean,
 })
 
-async def validate_input(hass, data: dict) -> dict:
+
+async def validate_input(hass, data: Dict[str, Any]) -> Dict[str, Any]:
     """Validate the user input allows us to connect."""
+    # Normalize host
+    host = data[CONF_HOST].strip()
+    if not host:
+        raise InvalidHost("Host cannot be empty")
+
+    # Remove protocol if provided
+    if host.startswith(("http://", "https://")):
+        host = host.split("://", 1)[1]
+    data[CONF_HOST] = host
+
+    # Validate port
+    port = data[CONF_PORT]
+    if not (1 <= port <= 65535):
+        raise InvalidPort("Port must be between 1 and 65535")
+
+    # Create session with appropriate SSL settings
     session = async_get_clientsession(hass, data.get(CONF_VERIFY_SSL, True))
 
+    # Create API instance
     api = AdGuardHomeAPI(
-        host=data[CONF_HOST],
-        port=data[CONF_PORT], 
+        host=host,
+        port=port,
         username=data.get(CONF_USERNAME),
         password=data.get(CONF_PASSWORD),
         ssl=data.get(CONF_SSL, False),
         session=session,
+        timeout=10,  # 10 second timeout for setup
     )
 
     # Test the connection
-    if not await api.test_connection():
-        raise CannotConnect
-
-    # Get server info
     try:
-        status = await api.get_status()
-        version = status.get("version", "unknown")
-        return {
-            "title": f"AdGuard Control Hub ({data[CONF_HOST]})",
-            "version": version
-        }
-    except Exception as err:
-        _LOGGER.exception("Unexpected exception: %s", err)
+        if not await api.test_connection():
+            raise CannotConnect("Failed to connect to AdGuard Home")
+
+        # Get additional server info if possible
+        try:
+            status = await api.get_status()
+            version = status.get("version", "unknown")
+            dns_port = status.get("dns_port", "N/A")
+
+            return {
+                "title": f"AdGuard Control Hub ({host})",
+                "version": version,
+                "dns_port": dns_port,
+                "host": host,
+            }
+        except Exception as err:
+            _LOGGER.warning("Could not get server status, but connection works: %s", err)
+            return {
+                "title": f"AdGuard Control Hub ({host})",
+                "version": "unknown",
+                "dns_port": "N/A",
+                "host": host,
+            }
+
+    except AdGuardAuthError as err:
+        _LOGGER.error("Authentication failed: %s", err)
+        raise InvalidAuth from err
+    except AdGuardConnectionError as err:
+        _LOGGER.error("Connection failed: %s", err)
+        if "timeout" in str(err).lower():
+            raise Timeout from err
         raise CannotConnect from err
+    except asyncio.TimeoutError as err:
+        _LOGGER.error("Connection timeout: %s", err)
+        raise Timeout from err
+    except Exception as err:
+        _LOGGER.exception("Unexpected error during validation: %s", err)
+        raise CannotConnect from err
+
 
 class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for AdGuard Control Hub."""
 
     VERSION = 1
+    CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_POLL
 
-    async def async_step_user(self, user_input: dict[str, Any] | None = None):
+    async def async_step_user(
+        self, user_input: Optional[Dict[str, Any]] = None
+    ) -> FlowResult:
         """Handle the initial step."""
-        errors: dict[str, str] = {}
+        errors: Dict[str, str] = {}
 
         if user_input is not None:
             try:
                 info = await validate_input(self.hass, user_input)
-            except CannotConnect:
-                errors["base"] = "cannot_connect"
-            except Exception:  
-                _LOGGER.exception("Unexpected exception")
-                errors["base"] = "unknown"
-            else:
+
                 # Create unique ID based on host and port
-                unique_id = f"{user_input[CONF_HOST]}:{user_input[CONF_PORT]}"
+                unique_id = f"{info['host']}:{user_input[CONF_PORT]}"
                 await self.async_set_unique_id(unique_id)
                 self._abort_if_unique_id_configured()
 
@@ -76,11 +131,83 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                     data=user_input,
                 )
 
+            except CannotConnect:
+                errors["base"] = "cannot_connect"
+            except InvalidAuth:
+                errors["base"] = "invalid_auth"
+            except InvalidHost:
+                errors[CONF_HOST] = "invalid_host"
+            except InvalidPort:
+                errors[CONF_PORT] = "invalid_port"
+            except Timeout:
+                errors["base"] = "timeout"
+            except Exception:
+                _LOGGER.exception("Unexpected exception during config flow")
+                errors["base"] = "unknown"
+
         return self.async_show_form(
             step_id="user",
             data_schema=STEP_USER_DATA_SCHEMA,
             errors=errors,
         )
 
+    async def async_step_import(self, import_info: Dict[str, Any]) -> FlowResult:
+        """Handle configuration import."""
+        return await self.async_step_user(import_info)
+
+    @staticmethod
+    def async_get_options_flow(config_entry):
+        """Get the options flow for this handler."""
+        return OptionsFlowHandler(config_entry)
+
+
+class OptionsFlowHandler(config_entries.OptionsFlow):
+    """Handle options flow for AdGuard Control Hub."""
+
+    def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
+        """Initialize options flow."""
+        self.config_entry = config_entry
+
+    async def async_step_init(
+        self, user_input: Optional[Dict[str, Any]] = None
+    ) -> FlowResult:
+        """Handle options flow."""
+        if user_input is not None:
+            return self.async_create_entry(title="", data=user_input)
+
+        options_schema = vol.Schema({
+            vol.Optional(
+                "scan_interval",
+                default=self.config_entry.options.get("scan_interval", 30),
+            ): vol.All(vol.Coerce(int), vol.Range(min=10, max=300)),
+            vol.Optional(
+                "timeout",
+                default=self.config_entry.options.get("timeout", 10),
+            ): vol.All(vol.Coerce(int), vol.Range(min=5, max=60)),
+        })
+
+        return self.async_show_form(
+            step_id="init",
+            data_schema=options_schema,
+        )
+
+
+# Custom exceptions
 class CannotConnect(Exception):
-    """Error to indicate we cannot connect."""
+    """Error to indicate we cannot connect."""
+
+
+class InvalidAuth(Exception):
+    """Error to indicate there is invalid auth."""
+
+
+class InvalidHost(Exception):
+    """Error to indicate invalid host."""
+
+
+class InvalidPort(Exception):
+    """Error to indicate invalid port."""
+
+
+class Timeout(Exception):
+    """Error to indicate connection timeout."""