Fix Octopus account validation, improve config flow error handling

custom_components/octopus_intelligent/config_flow.py

@@ -23,6 +23,7 @@
     CONF_OFFPEAK_END_DEFAULT,
     INTELLIGENT_24HR_TIMES
 )
+from .graphql_util import InvalidAuthError, validate_octopus_account
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,7 +33,7 @@ class OctopusIntelligentConfigFlowHandler(config_entries.ConfigFlow, domain=DOMA
 
     async def _show_setup_form(self, errors=None):
         """Show the setup form to the user."""
-        errors = {}
+        errors = errors or {}
 
         fields = OrderedDict()
         fields[vol.Required(CONF_API_KEY)] = str
@@ -59,9 +60,12 @@ async def async_step_user(self, user_input=None):
         errors = {}
         try:
             await try_connection(user_input[CONF_API_KEY], user_input[CONF_ACCOUNT_ID])
-        except Exception:  # pylint: disable=broad-except
-            _LOGGER.exception("Unexpected exception")
-            errors["base"] = "unknown"
+        except Exception as ex:  # pylint: disable=broad-except
+            _LOGGER.error(ex)
+            if isinstance(ex, InvalidAuthError):
+                errors["base"] = "invalid_auth"
+            else:
+                errors["base"] = "unknown"
             return await self._show_setup_form(errors)
 
         unique_id = user_input[CONF_ACCOUNT_ID]
@@ -118,18 +122,9 @@ async def async_step_user(self, user_input=None):
 #         return self.async_show_form(step_id="user", data_schema=vol.Schema(fields), errors=errors)
 
 
-async def try_connection(api_key, account_id):
+async def try_connection(api_key: str, account_id: str):
+    """Try connecting to the Octopus API and validating the given account_id."""
     _LOGGER.debug("Trying to connect to Octopus during setup")
     client = OctopusEnergyGraphQLClient(api_key)
-    try:
-        accounts = await client.async_get_accounts()
-        if (account_id not in accounts):
-            _LOGGER.error(f"Account {account_id} not found in accounts {accounts}")
-            raise Exception(f"Account {account_id} not found in accounts {accounts}")
-    except Exception as ex:
-        _LOGGER.error(f"Authentication failed : {ex.message}. You may need to check your token or create a new app in the gardena api and use the new token.")
-
-    # smart_system = SmartSystem(email=email, password=password, client_id=client_id)
-    # smart_system.authenticate()
-    # smart_system.update_locations()
+    await validate_octopus_account(client, account_id)
     _LOGGER.debug("Successfully connected to Octopus during setup")

custom_components/octopus_intelligent/graphql_client.py

@@ -17,7 +17,7 @@ def __init__(self, api_key: str):
     self._login_attempt = 0
     self._session = None
 
-  async def async_get_accounts(self):
+  async def async_get_accounts(self) -> list[str]:
     """Gets the accounts for the given API key"""
     return await self.__async_execute_with_session(self.__async_get_accounts)
 

custom_components/octopus_intelligent/graphql_util.py

@@ -0,0 +1,63 @@
+"""Validation and parsing for the Octopus GraphQL API."""
+from ast import literal_eval
+from pprint import pformat
+
+from gql.transport.exceptions import TransportQueryError
+
+from .graphql_client import OctopusEnergyGraphQLClient
+
+
+class InvalidAuthError(Exception):
+    """Invalid Octopus API key or account number."""
+
+
+async def validate_octopus_account(client: OctopusEnergyGraphQLClient, account_id: str):
+    """Check that the Octopus account_id matches the authenticated API."""
+    try:
+        accounts = await client.async_get_accounts()
+    except TransportQueryError as ex:
+        msg = parse_gql_query_error(ex, "Authentication failed")
+        raise InvalidAuthError(msg) from ex
+    if account_id not in accounts:
+        raise InvalidAuthError(
+            f"Account '{account_id}' not found in accounts {accounts}"
+        )
+
+
+def parse_gql_query_error(error: TransportQueryError, default_title: str) -> str:
+    """Format a GQL JSON-like error response into something arguably readable for logging.
+
+    Sample formatted return value:
+
+    Authentication failed:
+    {'message': 'Invalid data.',
+     'path': ['obtainKrakenToken'],
+     'extensions': {'errorType': 'VALIDATION',
+                    'errorCode': 'KT-CT-1139',
+                    'errorDescription': 'Authentication failed.',
+                    'errorClass': 'VALIDATION',
+                    'validationErrors': [{'message': 'Authentication failed.',
+                                          'inputPath': ['input', 'apiKey']}]}}
+    """
+    err_str = str(error)
+    if len(err_str) > 500:  # Some protection against wild inputs
+        return err_str
+    try:
+        # Using ast.literal_eval() instead of json.loads() because the GQL
+        # response uses single quotes in object notation, and quote replacement
+        # may go wrong if the error messages include quotes.
+        obj = literal_eval(err_str)
+        if not isinstance(obj, dict):
+            return err_str
+        if "locations" in obj:
+            del obj["locations"]  # Noise
+        exts = obj.get("extensions", {})
+        title = exts.get("errorDescription", "") if isinstance(exts, dict) else ""
+        if isinstance(title, str) and len(title) > 1:
+            title = title.rstrip(".") or default_title
+        else:
+            title = default_title
+        body = pformat(obj, sort_dicts=False)
+        return f"{title}:\n{body}"
+    except Exception:  # pylint: disable=broad-except
+        return err_str

custom_components/octopus_intelligent/manifest.json

@@ -7,5 +7,5 @@
   "codeowners": ["@megakid"],
   "iot_class": "cloud_polling",
   "config_flow": true,
-  "requirements": ["gql==3.2.0"]
+  "requirements": ["gql~=3.5.0"]
 }

custom_components/octopus_intelligent/octopus_intelligent_system.py

@@ -11,6 +11,7 @@
 )
 
 from .graphql_client import OctopusEnergyGraphQLClient
+from .graphql_util import validate_octopus_account
 from .util import *
 
 _LOGGER = logging.getLogger(__name__)
@@ -173,13 +174,7 @@ async def async_cancel_boost_charge(self):
 
     async def start(self):
         _LOGGER.debug("Starting OctopusIntelligentSystem")
-        try:
-            accounts = await self.client.async_get_accounts()
-            if (self._account_id not in accounts):
-                _LOGGER.error(f"Account {self._account_id} not found in accounts {accounts}")
-                raise Exception(f"Account {self._account_id} not found in accounts {accounts}")
-        except Exception as ex:
-            _LOGGER.error(f"Authentication failed : {ex.message}. You may need to check your token or create a new app in the gardena api and use the new token.")
+        await validate_octopus_account(self.client, self._account_id)
 
     async def stop(self):
         _LOGGER.debug("Stopping OctopusIntelligentSystem")