add multiple contracts

duhow · Feb 8, 2024 · e73285b · e73285b
1 parent 17bac09
commit e73285b
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 27 deletions.
diff --git a/README.md b/README.md
@@ -33,8 +33,8 @@ Cuando lo tengas descargado, agrega la integración.
 ## To-Do
 
 - [x] Sensor de último consumo disponible
+- [x] Soportar múltiples contratos
 - [ ] Publicar el consumo en [Energía](https://www.home-assistant.io/docs/energy/)
-- [ ] Soportar múltiples contratos **(?)**
 
 ## Ayuda
 

diff --git a/custom_components/aigues_barcelona/config_flow.py b/custom_components/aigues_barcelona/config_flow.py
@@ -71,10 +71,8 @@ async def validate_credentials(
             _LOGGER.info("Login succeeded!")
         contracts = await hass.async_add_executor_job(api.contracts, username)
 
-        if len(contracts) > 1:
-            raise NotImplementedError("Multiple contracts are not supported")
-        # FIXME calling function instead of property due to async incompatibility
-        return {CONF_CONTRACT: contracts[0]["contractDetail"]["contractNumber"]}
+        available_contracts = [x["contractDetail"]["contractNumber"] for x in contracts]
+        return {CONF_CONTRACT: available_contracts}
 
     except Exception:
         _LOGGER.debug(f"Last data: {api.last_response}")
@@ -134,8 +132,8 @@ async def async_step_reauth_confirm(
         user_input = {**self.stored_input, **user_input}
         try:
             info = await validate_credentials(self.hass, user_input)
-            contract = info[CONF_CONTRACT]
-            if contract != self.stored_input.get(CONF_CONTRACT):
+            contracts = info[CONF_CONTRACT]
+            if contracts != self.stored_input.get(CONF_CONTRACT):
                 _LOGGER.error("Reauth failed, contract does not match stored one")
                 raise InvalidAuth
 
@@ -172,9 +170,9 @@ async def async_step_user(
             _LOGGER.debug(f"Result is {info}")
             if not info:
                 raise InvalidAuth
-            contract = info[CONF_CONTRACT]
+            contracts = info[CONF_CONTRACT]
 
-            await self.async_set_unique_id(contract.lower())
+            await self.async_set_unique_id(user_input["username"])
             self._abort_if_unique_id_configured()
         except NotImplementedError:
             errors["base"] = "not_implemented"
@@ -193,10 +191,11 @@ async def async_step_user(
         except AlreadyConfigured:
             errors["base"] = "already_configured"
         else:
-            _LOGGER.debug(f"Creating entity with {user_input} and {contract=}")
+            _LOGGER.debug(f"Creating entity with {user_input} and {contracts=}")
+            nif_oculto = user_input[CONF_USERNAME][-3:][0:2]
 
             return self.async_create_entry(
-                title=f"Aigua {contract}", data={**user_input, **info}
+                title=f"Aigua ****{nif_oculto}", data={**user_input, **info}
             )
 
         return self.async_show_form(

diff --git a/custom_components/aigues_barcelona/sensor.py b/custom_components/aigues_barcelona/sensor.py
@@ -24,6 +24,7 @@
 from .const import ATTR_LAST_MEASURE
 from .const import CONF_CONTRACT
 from .const import CONF_VALUE
+from .const import DEFAULT_SCAN_PERIOD
 from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
@@ -37,27 +38,32 @@ async def async_setup_entry(hass: HomeAssistant, config_entry, async_add_entitie
 
     username = config_entry.data[CONF_USERNAME]
     password = config_entry.data[CONF_PASSWORD]
-    contract = config_entry.data[CONF_CONTRACT]
+    contracts = config_entry.data[CONF_CONTRACT]
     token = config_entry.data.get(CONF_TOKEN)
 
-    coordinator = ContratoAgua(
-        hass, username, password, contract, token=token, prev_data=None
-    )
+    contadores = list()
 
-    # postpone first refresh to speed up startup
-    @callback
-    async def async_first_refresh(*args):
-        """Force the component to assess the first refresh."""
-        await coordinator.async_refresh()
+    for contract in contracts:
+        coordinator = ContratoAgua(
+            hass, username, password, contract, token=token, prev_data=None
+        )
+
+        # postpone first refresh to speed up startup
+        @callback
+        async def async_first_refresh(*args):
+            """Force the component to assess the first refresh."""
+            await coordinator.async_refresh()
+
+        if hass.state == CoreState.running:
+            await async_first_refresh()
+        else:
+            hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START, async_first_refresh)
 
-    if hass.state == CoreState.running:
-        await async_first_refresh()
-    else:
-        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_START, async_first_refresh)
+        contadores.append(ContadorAgua(coordinator))
 
     _LOGGER.info("about to add entities")
     # add sensor entities
-    async_add_entities([ContadorAgua(coordinator)])
+    async_add_entities(contadores)
 
     return True
 
@@ -94,7 +100,7 @@ def __init__(
             hass,
             _LOGGER,
             name=self.id,
-            update_interval=timedelta(minutes=240),
+            update_interval=timedelta(seconds=DEFAULT_SCAN_PERIOD),
         )
 
     async def _async_update_data(self):
@@ -120,7 +126,7 @@ async def _async_update_data(self):
             # TODO: change once recaptcha is fiexd
             # await self.hass.async_add_executor_job(self._api.login)
             consumptions = await self.hass.async_add_executor_job(
-                self._api.consumptions, LAST_WEEK, TOMORROW
+                self._api.consumptions, LAST_WEEK, TOMORROW, self.contract
             )
         except ConfigEntryAuthFailed as exp:
             _LOGGER.error("Token has expired, cannot check consumptions.")