From 6b2368300714773b62e6dbb521304052daf49512 Mon Sep 17 00:00:00 2001
From: Daniel Vejsada <dan.vejsada@gmail.com>
Date: Thu, 15 Feb 2024 22:20:20 +0100
Subject: [PATCH] Improving error handling during config.

---
 pid_integration/api_call.py    |  8 ++----
 pid_integration/config_flow.py | 52 ++++++++++++++++++++++++++--------
 2 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/pid_integration/api_call.py b/pid_integration/api_call.py
index 5d0b4f2..e87bad3 100644
--- a/pid_integration/api_call.py
+++ b/pid_integration/api_call.py
@@ -10,21 +10,17 @@ class ApiCall:
     def update_info(api_key, stop_id, conn_num):
         headers = {"Content-Type": "application/json; charset=utf-8", "x-access-token": api_key}
         parameters = {"aswIds": stop_id, "total": conn_num, "minutesAfter": 10080}
-
         response = requests.get(API_URL, params=parameters, headers=headers)
         return response.json()
 
+
     @staticmethod
     def authenticate(api_key, stop_id, conn_num):
         headers = {"Content-Type": "application/json; charset=utf-8", "x-access-token": api_key}
         parameters = {"aswIds": stop_id, "total": conn_num, "minutesAfter": 10080}
         response = requests.get(API_URL, params=parameters, headers=headers)
         reply = response.json()
-        title = reply["stops"][0]["stop_name"] + " " + ApiCall.check_not_null(reply["stops"][0]["platform_code"])
-        if response.status_code == 200:
-            return True, title
-        else:
-            return False, None
+        return response.status_code, reply
 
     @staticmethod
     def check_not_null(response):
diff --git a/pid_integration/config_flow.py b/pid_integration/config_flow.py
index 45021c3..b99de04 100644
--- a/pid_integration/config_flow.py
+++ b/pid_integration/config_flow.py
@@ -18,14 +18,17 @@ async def validate_input(hass: HomeAssistant, data: dict) -> dict[str, Any]:
     Data has the keys from DATA_SCHEMA with values provided by the user.
     """
 
-    result, stop = await hass.async_add_executor_job(ApiCall.authenticate, data[CONF_API_KEY], data[CONF_ID], data[CONF_DEP_NUM])
-
-    if not result:
-
+    status, reply = await hass.async_add_executor_job(ApiCall.authenticate, data[CONF_API_KEY], data[CONF_ID], data[CONF_DEP_NUM])
+    if status == 200:
+        title: str = reply["stops"][0]["stop_name"] + " " + ApiCall.check_not_null(reply["stops"][0]["platform_code"])
+        return {"title": title}
+    elif status == 401:
+        raise WrongApiKey
+    elif status == 404:
+        raise StopNotFound
+    else:
         raise CannotConnect
 
-    return {"title": stop}
-
 
 class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
@@ -34,27 +37,44 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     async def async_step_user(self, user_input=None):
 
-        # Check for any previous instance of the integration and preload the API key.
+        # Check for any previous instance of the integration
         if DOMAIN in list(self.hass.data.keys()):
+            # If previous instance exists, set the API key as suggestion to new config
             data_schema = vol.Schema(
                 {vol.Required(CONF_ID): str, vol.Required(CONF_API_KEY, default=self.hass.data[DOMAIN][list(self.hass.data[DOMAIN].keys())[0]].api_key): str,
                 vol.Required(CONF_DEP_NUM): int}
                 )
         else:
+            # if no previous instance, show blank form
             data_schema = vol.Schema(
                 {vol.Required(CONF_ID): str, vol.Required(CONF_API_KEY): str,
                  vol.Required(CONF_DEP_NUM): int}
             )
 
-        errors = {}
+        # Set dict for errors
+        errors: dict = {}
+
+        # Steps to take if user input is received
         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 CannotConnect:
+                _LOGGER.exception("Unknown API connection error")
+                errors["base"] = "Unknown API connection error"
+
+            except WrongApiKey:
+                _LOGGER.exception("The API key provided in configuration is not correct.")
+                errors["base"] = "Connection was not authorized. Wrong or no API key provided"
+
+            except StopNotFound:
+                _LOGGER.exception("Stop with provided awsIDs was not found.")
+                errors["base"] = "Non existent awsIds provided, stop not found."
+
             except Exception:  # pylint: disable=broad-except
-                _LOGGER.exception("Unexpected exception")
-                errors["base"] = "unknown"
+                _LOGGER.exception("Unknown exception - cannot connect to API")
+                errors["base"] = "Unknown exception - cannot connect to API"
 
         # 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(
@@ -63,4 +83,12 @@ async def async_step_user(self, user_input=None):
 
 
 class CannotConnect(exceptions.HomeAssistantError):
-    """Error to indicate we cannot connect."""
+    """Error to indicate we cannot connect for unknown reason."""
+
+
+class WrongApiKey(exceptions.HomeAssistantError):
+    """Error to indicate wrong API key was provided."""
+
+
+class StopNotFound(exceptions.HomeAssistantError):
+    """Error to indicate wrong stop was provided."""