From 5bfaae0d9a55f2342fee838bb9cdf98666fe4b1b Mon Sep 17 00:00:00 2001
From: Mathieu Broillet <mathieu@broillet.ch>
Date: Sat, 30 Dec 2023 21:22:32 +0100
Subject: [PATCH] code refactor

---
 .../easy_computer_manager/__init__.py         | 11 ++++----
 .../easy_computer_manager/config_flow.py      | 28 ++++++-------------
 2 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/custom_components/easy_computer_manager/__init__.py b/custom_components/easy_computer_manager/__init__.py
index 2d2bc26..923730b 100644
--- a/custom_components/easy_computer_manager/__init__.py
+++ b/custom_components/easy_computer_manager/__init__.py
@@ -1,9 +1,8 @@
 """The Easy Dualboot Computer Manager integration."""
 
-# Some code is from the official wake_on_lan integration
+# Some snippets of code are from the official wake_on_lan integration (inspiration for this custom component)
 
 from __future__ import annotations
-
 import logging
 from functools import partial
 
@@ -28,10 +27,10 @@ WAKE_ON_LAN_SEND_MAGIC_PACKET_SCHEMA = vol.Schema(
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    """Set up the wake on LAN component."""
+    """Set up the Easy Dualboot Computer Manager integration."""
 
     async def send_magic_packet(call: ServiceCall) -> None:
-        """Send magic packet to wake up a device."""
+        """Send a magic packet to wake up a device."""
         mac_address = call.data.get(CONF_MAC)
         broadcast_address = call.data.get(CONF_BROADCAST_ADDRESS)
         broadcast_port = call.data.get(CONF_BROADCAST_PORT)
@@ -43,7 +42,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             service_kwargs["port"] = broadcast_port
 
         _LOGGER.info(
-            "Send magic packet to mac %s (broadcast: %s, port: %s)",
+            "Sending magic packet to MAC %s (broadcast: %s, port: %s)",
             mac_address,
             broadcast_address,
             broadcast_port,
@@ -70,7 +69,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    """Unload a config entry."""
+    """Unload the Easy Dualboot Computer Manager integration."""
     return await hass.config_entries.async_forward_entry_unload(
         entry, "switch"
     )
diff --git a/custom_components/easy_computer_manager/config_flow.py b/custom_components/easy_computer_manager/config_flow.py
index 6a02776..0b5804f 100644
--- a/custom_components/easy_computer_manager/config_flow.py
+++ b/custom_components/easy_computer_manager/config_flow.py
@@ -7,6 +7,7 @@ from typing import Any
 import voluptuous as vol
 from homeassistant import config_entries, exceptions
 from homeassistant.core import HomeAssistant
+
 from paramiko.ssh_exception import AuthenticationException
 
 from . import utils
@@ -48,8 +49,7 @@ class Hub:
     async def test_connection(self) -> bool:
         """Test connectivity to the computer is OK."""
         try:
-            return utils.test_connection(
-                utils.create_ssh_connection(self._host, self._username, self._password, self._port))
+            return utils.test_connection(utils.create_ssh_connection(self._host, self._username, self._password, self._port))
         except AuthenticationException:
             return False
 
@@ -61,10 +61,9 @@ async def validate_input(hass: HomeAssistant, data: dict) -> dict[str, Any]:
     if len(data["host"]) < 3:
         raise InvalidHost
 
-    hub = Hub(hass, data["host"], data["username"], data["password"], data["port"])
+    hub = Hub(hass, **data)
 
-    result = await hub.test_connection()
-    if not result:
+    if not await hub.test_connection():
         raise CannotConnect
 
     return {"title": data["host"]}
@@ -81,23 +80,14 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         if user_input is not None:
             try:
                 info = await validate_input(self.hass, user_input)
-
                 return self.async_create_entry(title=info["title"], data=user_input)
-            except AuthenticationException:
-                errors["host"] = "cannot_connect"
-            except CannotConnect:
-                errors["base"] = "cannot_connect"
-            except InvalidHost:
-                errors["host"] = "cannot_connect"
-            except Exception:  # pylint: disable=broad-except
-                _LOGGER.exception("Unexpected exception")
+            except (AuthenticationException, CannotConnect, InvalidHost) as ex:
+                errors["base"] = str(ex)
+            except Exception as ex:  # pylint: disable=broad-except
+                _LOGGER.exception("Unexpected exception: %s", ex)
                 errors["base"] = "unknown"
 
-        # If there is no user input or there were errors, show the form again, including any errors that were found
-        # with the input.
-        return self.async_show_form(
-            step_id="user", data_schema=DATA_SCHEMA, errors=errors
-        )
+        return self.async_show_form(step_id="user", data_schema=DATA_SCHEMA, errors=errors)
 
 
 class CannotConnect(exceptions.HomeAssistantError):