fix: multiple contracts config flow fix (3) (#63)

* fix: multiple contracts config flow fix (3)

* chore: bump to 0.0.20b4

Author: GuyKh

custom_components/iec/config_flow.py

@@ -120,12 +120,12 @@ async def async_step_mfa(
                 data[CONF_BP_NUMBER] = customer.bp_number
 
                 contracts = await client.get_contracts(customer.bp_number)
-                contract_ids = [contract.contract_id for contract in contracts if contract.status == 1]
+                contract_ids = [int(contract.contract_id) for contract in contracts if contract.status == 1]
                 if len(contract_ids) == 1:
                     data[CONF_SELECTED_CONTRACTS] = [contract_ids[0]]
                     return self._async_create_iec_entry(data)
                 else:
-                    data[CONF_AVAILABLE_CONTRACTS] = [cid.lstrip('0') for cid in contract_ids]
+                    data[CONF_AVAILABLE_CONTRACTS] = contract_ids
                     self.data = data
                     return await self.async_step_select_contracts()
 

custom_components/iec/coordinator.py

@@ -36,7 +36,7 @@
 
 
 async def _verify_daily_readings_exist(daily_readings: list[RemoteReading], desired_date: datetime, device: Device,
-                                       contract_id: str, api: IecClient,
+                                       contract_id: int, api: IecClient,
                                        prefetched_reading: RemoteReadingResponse | None = None):
     desired_date = desired_date.replace(hour=0, minute=0, second=0, microsecond=0)
     daily_reading = next(filter(lambda x: find_reading_by_date(x, desired_date), daily_readings), None)
@@ -80,7 +80,7 @@ def __init__(
         )
         self._config_entry = config_entry
         self._bp_number = config_entry.data.get(CONF_BP_NUMBER)
-        self._contracts = config_entry.data.get(CONF_SELECTED_CONTRACTS)
+        self._contract_ids = config_entry.data.get(CONF_SELECTED_CONTRACTS)
         self._entry_data = config_entry.data
         self._today_readings = {}
         self.api = IecClient(
@@ -127,28 +127,32 @@ async def _async_update_data(
             self._bp_number = customer.bp_number
 
         all_contracts: list[Contract] = await self.api.get_contracts(self._bp_number)
-        if not self._contracts:
-            self._contracts = [contract.contract_id for contract in all_contracts if contract.status == 1]
+        if not self._contract_ids:
+            self._contract_ids = [int(contract.contract_id) for contract in all_contracts if contract.status == 1]
 
-        contracts: dict[str, Contract] = {c.contract_id: c for c in all_contracts}
+        contracts: dict[int, Contract] = {int(c.contract_id): c for c in all_contracts if c.status == 1
+                                          and int(c.contract_id) in self._contract_ids}
 
         tariff = await self.api.get_kwh_tariff() / 100
         data = {STATICS_DICT_NAME: {
             STATIC_KWH_TARIFF: tariff,
             STATIC_BP_NUMBER: self._bp_number
         }}
 
-        for contract_id in self._contracts:
+        _LOGGER.debug(f"All Contract Ids: {list(contracts.keys())}")
+
+        for contract_id in self._contract_ids:
             # Because IEC API provides historical usage/cost with a delay of a couple of days
             # we need to insert data into statistics.
+            _LOGGER.debug(f"Processing {contract_id}")
             await self._insert_statistics(contract_id, contracts.get(contract_id).smart_meter)
             billing_invoices = await self.api.get_billing_invoices(self._bp_number, contract_id)
             billing_invoices.invoices.sort(key=lambda inv: inv.full_date, reverse=True)
             last_invoice = billing_invoices.invoices[0]
 
             future_consumption: FutureConsumptionInfo | None = None
             daily_readings: list[RemoteReading] | None = None
-            today_reading: RemoteReadingResponse | None = None
+
             if contracts.get(contract_id).smart_meter:
                 # For some reason, there are differences between sending 2024-03-01 and sending 2024-03-07 (Today)
                 # So instead of sending the 1st day of the month, just sending today date
@@ -219,25 +223,27 @@ async def _async_update_data(
                             else:
                                 _LOGGER.debug("Failed fetching FutureConsumption, data in IEC API is corrupted")
 
-                data[contract_id] = {CONTRACT_DICT_NAME: contracts.get(contract_id),
-                                     INVOICE_DICT_NAME: last_invoice,
-                                     FUTURE_CONSUMPTIONS_DICT_NAME: future_consumption,
-                                     DAILY_READINGS_DICT_NAME: daily_readings,
-                                     STATICS_DICT_NAME: {STATIC_KWH_TARIFF: tariff} # workaround
-                                     }
+            data[str(contract_id)] = {CONTRACT_DICT_NAME: contracts.get(contract_id),
+                                      INVOICE_DICT_NAME: last_invoice,
+                                      FUTURE_CONSUMPTIONS_DICT_NAME: future_consumption,
+                                      DAILY_READINGS_DICT_NAME: daily_readings,
+                                      STATICS_DICT_NAME: {STATIC_KWH_TARIFF: tariff}  # workaround
+                                      }
+
+        # Clean today reading for next reading cycle
+        self._today_readings = {}
 
-            # Clean today reading for next reading cycle
-            self._today_readings = {}
+        _LOGGER.debug(f"Data Keys: {list(data.keys())}")
         return data
 
-    async def _insert_statistics(self, contract_id: str, is_smart_meter: bool) -> None:
+    async def _insert_statistics(self, contract_id: int, is_smart_meter: bool) -> None:
         if not is_smart_meter:
             _LOGGER.info(f"IEC Contract {contract_id} doesn't contain Smart Meters, not adding statistics")
             # Support only smart meters at the moment
             return
 
         _LOGGER.debug(f"Updating statistics for IEC Contract {contract_id}")
-        devices = await self.api.get_devices(contract_id)
+        devices = await self.api.get_devices(str(contract_id))
         month_ago_time = (datetime.now() - timedelta(weeks=4))
 
         kwh_price = await self.api.get_kwh_tariff() / 100
@@ -258,7 +264,7 @@ async def _insert_statistics(self, contract_id: str, is_smart_meter: bool) -> No
                 readings = await self.api.get_remote_reading(device.device_number, int(device.device_code),
                                                              month_ago_time,
                                                              month_ago_time, ReadingResolution.DAILY,
-                                                             contract_id)
+                                                             str(contract_id))
             else:
                 last_stat_time = last_stat[consumption_statistic_id][0]["start"]
                 # API returns daily data, so need to increase the start date by 4 hrs to get the next day
@@ -277,7 +283,7 @@ async def _insert_statistics(self, contract_id: str, is_smart_meter: bool) -> No
                 _LOGGER.debug(f"Fetching consumption from {from_date.strftime('%Y-%m-%d %H:%M:%S')}")
                 readings = await self.api.get_remote_reading(device.device_number, int(device.device_code),
                                                              from_date, from_date,
-                                                             ReadingResolution.DAILY, contract_id)
+                                                             ReadingResolution.DAILY, str(contract_id))
                 if from_date.date() == today.date():
                     self._today_readings[contract_id] = readings
 

custom_components/iec/manifest.json

@@ -10,5 +10,5 @@
   "iot_class": "cloud_polling",
   "issue_tracker": "https://github.com/guykh/iec-custom-component/issues",
   "requirements": ["iec-api==0.2.8"],
-  "version": "0.0.20b3"
+  "version": "0.0.20b4"
 }

custom_components/iec/sensor.py

@@ -255,7 +255,13 @@ def __init__(
     def native_value(self) -> StateType:
         """Return the state."""
         if self.coordinator.data is not None:
+            if self.contract_id == STATICS_DICT_NAME:
+                return self.entity_description.value_fn(
+                    self.coordinator.data.get(self.contract_id)
+                )
+
+            # Trim leading 0000 if needed and align with coordinator keys
             return self.entity_description.value_fn(
-                self.coordinator.data.get(self.contract_id)
+                self.coordinator.data.get(str(int(self.contract_id)))
             )
         return None