Adding config flow setup, convert async to sync

.gitignore

@@ -158,3 +158,7 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Home Assistant / Testing
+
+config/

custom_components/aigues_barcelona/__init__.py

@@ -1,34 +1,43 @@
-"""Example integration."""
+""" Integration for Aigues de Barcelona. """
 
-import asyncio
-from homeassistant.core import HomeAssistant
-from homeassistant.helpers.typing import ConfigType
+from __future__ import annotations
+
+import voluptuous as vol
+
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
-    CONF_NAME,
+    CONF_USERNAME,
     CONF_PASSWORD,
-    CONF_SCAN_INTERVAL
+    Platform,
 )
-import logging
-import voluptuous as vol
-import homeassistant.helpers.config_validation as cv
+from homeassistant.core import HomeAssistant, ServiceCall
+from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN, DEFAULT_SCAN_PERIOD
+from .const import DOMAIN
+from .api import AiguesApiClient
 
-_LOGGER = logging.getLogger(__name__)
+PLATFORMS = [Platform.SENSOR]
 
-ACCOUNT_CONFIG_SCHEMA = vol.Schema(
-    {
-        vol.Required(CONF_NAME): cv.string,
-        vol.Required(CONF_PASSWORD): cv.string,
-        vol.Optional(CONF_SCAN_INTERVAL, default=DEFAULT_SCAN_PERIOD): cv.time_period,
-    }
-)
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+    api = AiguesApiClient(entry.data[CONF_USERNAME], entry.data[CONF_PASSWORD])
+
+    try:
+        await hass.async_add_executor_job(api.login)
+    except:
+        raise ConfigEntryNotReady
+
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+
+    return True
 
-async def async_setup(hass: HomeAssistant, config: ConfigType, discovery_info=None) -> bool:
-    """ Set up the domain. """
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    if unload_ok:
+        hass.data[DOMAIN].pop(entry.entry_id)
+    if not hass.data[DOMAIN]:
+        del hass.data[DOMAIN]
 
-    if DOMAIN not in config:
-        _LOGGER.debug(
-            "Nothing to import from configuration.yaml, loading from Integrations",
-        )
-        return True
+    return unload_ok

custom_components/aigues_barcelona/api.py

@@ -1,17 +1,16 @@
-import asyncio
 import logging
-import socket
+import requests
 
-import aiohttp
-import async_timeout
 import datetime
 
-TIMEOUT = 25
+TIMEOUT = 60
 
-_LOGGER: logging.Logger = logging.getLogger(__package__)
+_LOGGER: logging.Logger = logging.getLogger(__name__)
 
 class AiguesApiClient:
-    def __init__(self, session: aiohttp.ClientSession, username, password, contract=None):
+    def __init__(self, username, password, contract=None, session: requests.Session = None):
+        if session is None:
+            session = requests.Session()
         self.cli = session
         self.api_host = "https://api.aiguesdebarcelona.cat"
         # https://www.aiguesdebarcelona.cat/o/ofex-theme/js/chunk-vendors.e5935b72.js
@@ -33,7 +32,9 @@ def _generate_url(self, path, query) -> str:
 
     def _return_token_field(self, key):
         token = self.cli.cookies.get_dict().get("ofexTokenJwt")
-        assert token, "Token login missing"
+        if not token:
+            _LOGGER.warning("Token login missing")
+            return False
 
         data = token.split(".")[1]
         logging.debug(data)
@@ -42,28 +43,31 @@ def _return_token_field(self, key):
 
         return json.loads(data).get(key)
 
-    async def _query(self, path, query=None, json=None, headers=None, method="GET"):
+    def _query(self, path, query=None, json=None, headers=None, method="GET"):
         if headers is None:
             headers = dict()
         headers = {**self.headers, **headers}
 
-        async with self.cli.request(
+        resp = self.cli.request(
             method=method,
             url=self._generate_url(path, query),
             json=json,
-            headers=headers
-        ) as resp:
-            if resp.status == 404:
-                msg = resp.text()
-                if len(msg) > 5:
-                    msg = resp.json().get("message", r.text)
-                raise Exception(f"Not found: {msg}")
-            if resp.status == 401:
-                msg = resp.text()
-                if len(msg) > 5:
-                    msg = resp.json().get("message", r.text)
-                raise Exception(f"Denied: {msg}")
-            return await resp
+            headers=headers,
+            timeout=TIMEOUT
+        )
+        _LOGGER.debug(f"Query done with code {resp.status}")
+        msg = resp.text
+        if len(msg) > 5:
+            msg = resp.json().get("message", r.text)
+
+        if resp.status == 500:
+            raise Exception(f"Server error: {msg}")
+        if resp.status == 404:
+            raise Exception(f"Not found: {msg}")
+        if resp.status == 401:
+            raise Exception(f"Denied: {msg}")
+
+        return resp
 
     def login(self, user=None, password=None, recaptcha=None):
         if user is None:
@@ -90,10 +94,17 @@ def login(self, user=None, password=None, recaptcha=None):
 
         r = self._query(path, query, body, headers, method="POST")
 
-        assert not r.json().get("errorMessage"), r.json().get("errorMessage")
-        #{"errorMessage":"Los datos son incorrectos","result":false,"errorCode":"LOGIN_ERROR"}
-        access_token = r.json().get("access_token")
-        assert access_token, "Access token missing"
+        error = r.json().get("errorMessage", None)
+        if error:
+            _LOGGER.warning(error)
+            return False
+
+        access_token = r.json().get("access_token", None)
+        if not access_token:
+            _LOGGER.warning("Access token missing")
+            return False
+
+        return True
 
         # set as cookie: ofexTokenJwt
         # https://www.aiguesdebarcelona.cat/ca/area-clientes

custom_components/aigues_barcelona/config_flow.py

@@ -5,76 +5,102 @@
 from typing import Any
 
 import voluptuous as vol
-from homeassistant import config_entries, core, exceptions
-from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant import config_entries, core
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.const import (
+    CONF_USERNAME,
+    CONF_PASSWORD,
+    CONF_SCAN_INTERVAL
+)
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
+import homeassistant.helpers.config_validation as cv
 
-
-from . import const
+from .const import DOMAIN
 from .api import AiguesApiClient
 
 _LOGGER = logging.getLogger(__name__)
 
-STEP_USER_DATA_SCHEMA = vol.Schema(
+ACCOUNT_CONFIG_SCHEMA = vol.Schema(
     {
-        vol.Required(CONF_USERNAME): str,
-        vol.Required(CONF_PASSWORD): str,
+        vol.Required(CONF_USERNAME): cv.string,
+        vol.Required(CONF_PASSWORD): cv.string,
+        #vol.Optional(CONF_SCAN_INTERVAL, default=const.DEFAULT_SCAN_PERIOD): cv.time_period_seconds
     }
 )
 
+async def validate_credentials(hass: HomeAssistant, data: dict[str, Any]) -> dict[str, Any]:
+    username = data[CONF_USERNAME]
+    password = data[CONF_PASSWORD]
+
+    # QUICK check DNI/NIF. TODO improve
+    if len(username) != 9 or not username[0:8].isnumeric():
+        raise InvalidUsername
 
-async def _test_login_user(hass: HomeAssistant, user_input: dict[str, Any]) -> dict[str, Any]:
-    """Return true if credentials are valid"""
     try:
-        session = async_create_clientsession(hass)
-        client = AiguesApiClient(session, username, password)
-        await client.login()
-        return True
-    except Exception:
-        pass
-    return False
+        api = AiguesApiClient(username, password)
+        _LOGGER.info("Attempting to login")
+        login = await hass.async_add_executor_job(api.login)
+        if not login:
+            raise InvalidAuth
+        contracts = api.contract_id
+
+        if len(contracts) > 1:
+            raise NotImplementedError("Multiple contracts are not supported")
+        return {"contract": contracts[0]}
 
+    except Exception:
+        return False
 
-class ConfigFlow(config_entries.ConfigFlow, domain=const.DOMAIN):
-    """Handle a config flow."""
+class AiguesBarcelonaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
     async def async_step_user(
-        self, user_input: dict[str, Any] | None = None
+        self,
+        user_input: dict[str, Any] | None = None
     ) -> FlowResult:
-        """Handle the initial step."""
+        """Handle configuration step from UI."""
         if user_input is None:
             return self.async_show_form(
-                step_id="user", data_schema=STEP_USER_DATA_SCHEMA
+                step_id="user", data_schema=ACCOUNT_CONFIG_SCHEMA
             )
 
         errors = {}
 
         try:
-            info = await _test_login_user(self.hass, user_input)
+            info = await validate_credentials(self.hass, user_input)
+            _LOGGER.debug(f"Result is {info}")
+            if not info:
+                raise InvalidAuth
+            contract = info["contract"]
+
+            await self.async_set_unique_id(contract.lower())
+            self._abort_if_unique_id_configured()
+        except NotImplementedError:
+            errors["base"] = "not_implemented"
+        except InvalidUsername:
+            errors["base"] = "invalid_auth"
+        except InvalidAuth:
+            errors["base"] = "invalid_auth"
         except AlreadyConfigured:
             errors["base"] = "already_configured"
         else:
-            #await self.async_set_unique_id(user_input[const.CONF_CUPS])
-            #self._abort_if_unique_id_configured()
-            #extra_data = {"scups": user_input[const.CONF_CUPS][-4:]}
+            _LOGGER.debug(f"Creating entity with {user_input} and {contract=}")
+
             return self.async_create_entry(
-                title="test aigua", data={**user_input}
+                title=f"Aigua {contract}", data={**user_input, **info}
             )
 
         return self.async_show_form(
-            step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
+            step_id="user", data_schema=ACCOUNT_CONFIG_SCHEMA, errors=errors
         )
 
-    async def async_step_import(self, import_data: dict[str, Any]) -> FlowResult:
-        """Import data from yaml config"""
-        await self.async_set_unique_id(import_data[const.CONF_CUPS])
-        self._abort_if_unique_id_configured()
-        scups = import_data[const.CONF_CUPS][-4:]
-        extra_data = {"scups": scups}
-        return self.async_create_entry(title=scups, data={**import_data, **extra_data})
-
-class AlreadyConfigured(exceptions.HomeAssistantError):
+class AlreadyConfigured(HomeAssistantError):
     """Error to indicate integration is already configured."""
+
+class InvalidAuth(HomeAssistantError):
+    """Error to indicate credentials are invalid"""
+
+class InvalidUsername(HomeAssistantError):
+    """Error to indicate invalid username"""

custom_components/aigues_barcelona/const.py

@@ -1,6 +1,5 @@
 """ Constants definition """
-from datetime import timedelta
 
 DOMAIN = "aigues_barcelona"
 
-DEFAULT_SCAN_PERIOD = timedelta(hours=1)
+DEFAULT_SCAN_PERIOD = 3600

custom_components/aigues_barcelona/manifest.json

@@ -4,7 +4,7 @@
   "version": "0.0.1",
   "config_flow": true,
   "integration_type": "hub",
-  "documentation": "https://github.com/duhow/hass-aigues-barcelona/tree/main",
+  "documentation": "https://github.com/duhow/hass-aigues-barcelona/tree/master",
   "issue_tracker": "https://github.com/duhow/hass-aigues-barcelona/issues",
   "dependencies": ["http"],
   "codeowners": ["@duhow"],

custom_components/aigues_barcelona/sensor.py

@@ -3,25 +3,36 @@
 
 from homeassistant.core import HomeAssistant
 from homeassistant.components.sensor import SensorEntity
-from homeassistant.const.UnitOfVolume import CUBIC_METERS
+from homeassistant.const import UnitOfVolume
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 from . import DOMAIN
 
-
-def setup_platform(
+async def async_setup_platform(
     hass: HomeAssistant,
     config: ConfigType,
     add_entities: AddEntitiesCallback,
     discovery_info: DiscoveryInfoType | None = None
 ) -> None:
     """Set up the sensor platform."""
-    # We only want this platform to be set up via discovery.
-    if discovery_info is None:
+
+    username = config[CONF_USERNAME]
+    password = config[CONF_PASSWORD]
+    contract = config["contract"]
+    try:
+        client = AiguesApiClient(username, password, contract=contract)
+
+        login = await hass.async_add_executor_job(client.login)
+        if not login:
+            _LOGGER.warning("Wrong username and/or password")
+            return
+
+    except Exception:
+        _LOGGER.warning("Unable to create Aigues Barcelona Client")
         return
-    add_entities([ContadorAgua()])
 
+    async_add_entities([ContadorAgua(client)])
 
 class ContadorAgua(SensorEntity):
     """Representation of a sensor."""
@@ -43,7 +54,7 @@ def state(self):
     @property
     def unit_of_measurement(self) -> str:
         """Return the unit of measurement."""
-        return CUBIC_METERS
+        return UnitOfVolume.CUBIC_METERS
 
     def update(self) -> None:
         """Fetch new state data for the sensor.