Handle multiple contracts per customers

pyhydroquebec/__main__.py

@@ -14,14 +14,17 @@
 from pyhydroquebec.__version__ import VERSION
 
 
-async def fetch_data(client, contract_id, fetch_hourly=False):
+async def fetch_data(client, contract_id=None, fetch_hourly=False):
     """Fetch data for basic report."""
     await client.login()
     for customer in client.customers:
-        if customer.contract_id != contract_id and contract_id is not None:
+        if contract_id is None and customer.contract_list:
+            client.logger.warning("Contract id not specified, using first available: " + customer.selected_contract)
+        elif contract_id not in customer.contract_list:
             continue
-        if contract_id is None:
-            client.logger.warn("Contract id not specified, using first available.")
+        else:
+            customer.selected_contract = contract_id
+
 
         await customer.fetch_current_period()
         await customer.fetch_annual_data()
@@ -36,13 +39,16 @@ async def fetch_data(client, contract_id, fetch_hourly=False):
         if fetch_hourly:
             await customer.fetch_hourly_data(yesterday_str)
         return customer
-
+
+    client.logger.error("Could not fetch data, check contract id.")
+    return None
 
 async def dump_data(client, contract_id):
     """Fetch all data and dump them for debug and dev."""
     customer = await fetch_data(client, contract_id)
-    await customer.fetch_daily_data()
-    await customer.fetch_hourly_data()
+    if customer is not None:
+        await customer.fetch_daily_data()
+        await customer.fetch_hourly_data()
     return customer
 
 
@@ -51,7 +57,7 @@ async def list_contracts(client):
     await client.login()
     return [{"account_id": c.account_id,
              "customer_id": c.customer_id,
-             "contract_id": c.contract_id}
+             "contract_list": c.contract_list}
             for c in client.customers]
 
 
@@ -152,12 +158,16 @@ def main():
         loop.run_until_complete(close_fut)
         loop.close()
 
+    if results[0] is None:
+        return 1
+
     # Output data
     if args.list_contracts:
         for customer in results[0]:
-            print("Contract: {contract_id}\n\t"
-                  "Account: {account_id}\n\t"
-                  "Customer: {customer_id}".format(**customer))
+            print(f"Customer: {customer['customer_id']}\n\t"
+                  f"Account: {customer['account_id']}\n\t"
+                  f"Contract list: {', '.join(customer['contract_list'])}")
+
     elif args.dump_data:
         pprint(results[0].__dict__)
     elif args.influxdb:

pyhydroquebec/client.py

@@ -241,7 +241,7 @@ async def login(self):
             customer = Customer(self, account_id, customer_id, self._timeout, customer_logger)
             self._customers.append(customer)
             await customer.fetch_summary()
-            if customer.contract_id is None:
+            if not customer.contract_list:
                 del self._customers[-1]
 
     @property

pyhydroquebec/consts.py

@@ -126,12 +126,12 @@
 OVERVIEW_TPL = ("""
 ##################################
 # Hydro Quebec data for contract #
-# {0.contract_id}
+# {0.selected_contract}
 ##################################
 
 Account ID: {0.account_id}
 Customer number: {0.customer_id}
-Contract: {0.contract_id}
+Contract: {0.selected_contract}
 ===================
 
 Balance: {0.balance:.2f} $

pyhydroquebec/customer.py

@@ -19,15 +19,16 @@ class Customer():
 
     The account_id is called 'noPartenaireDemandeur' in the HydroQuebec API
     The customer_id is called 'Customer number' in the HydroQuebec 'My accounts' UI
-    The contract_id is called 'Contract' in the HydroQuebec 'At a glance' UI
+    The contract_list is called 'Contract' in the HydroQuebec 'At a glance' UI, several contracts might be listed
     """
 
     def __init__(self, client, account_id, customer_id, timeout, logger):
         """Create new Customer object."""
         self._client = client
         self.account_id = account_id
         self.customer_id = customer_id
-        self.contract_id = None
+        self._contract_list = []
+        self._selected_contract = None
         self._timeout = timeout
         self._logger = logger.getChild(customer_id)
         self._balance = None
@@ -57,19 +58,41 @@ async def fetch_summary(self):
             self._balance = float(raw_balance[:-2].replace(",", ".").
                                   replace("\xa0", ""))
 
-            raw_contract_id = soup.find('div', {'class': 'contrat'}).text
-            self.contract_id = (raw_contract_id
-                                .split("Contrat", 1)[-1]
-                                .replace("\t", "")
-                                .replace("\n", ""))
-
+            raw_contract_list = soup.find_all('div', {'class': 'contrat'}, limit=None)
+            for raw_contract_id in raw_contract_list:
+                self._contract_list.append(raw_contract_id.text
+                                    .split("Contrat", 1)[-1]
+                                    .replace("\t", "")
+                                    .replace("\n", ""))
+            if not self._contract_list:
+                self._logger.info("Customer has no contract")
+            else:
+                self._selected_contract = self._contract_list[0] # Select first contract by default
         except AttributeError:
-            self._logger.info("Customer has no contract")
+            self._logger.info("Error : a parsing error occured")
 
         # Needs to load the consumption profile page to not break
         # the next loading of the other pages
         await self._client.http_request(CONTRACT_CURRENT_URL_1, "get")
 
+    @property
+    def contract_list(self):
+        """Return the list of available contracts."""
+        return self._contract_list
+
+    @property
+    def selected_contract(self):
+        """Return the selected contract."""
+        return self._selected_contract
+
+    @selected_contract.setter
+    def selected_contract(self, contract_id):
+        """Select a specific contract."""
+        if contract_id not in self._contract_list:
+            raise ValueError(f"Contract {contract_id} not available. Possible contract are {', '.join(self._contract_list)}")
+        self._selected_contract = contract_id
+        self._logger.info("Contract %s selected", contract_id)
+
     @property
     def balance(self):
         """Return the collected balance."""
@@ -83,8 +106,9 @@ async def fetch_current_period(self):
         """
         self._logger.info("Fetching current period data")
         await self._client.select_customer(self.account_id, self.customer_id)
-
-        await self._client.http_request(CONTRACT_CURRENT_URL_1, "get")
+
+        params = {"noContrat": self._selected_contract}
+        await self._client.http_request(CONTRACT_CURRENT_URL_1, "get", params=params)
 
         headers = {"Content-Type": "application/json"}
         res = await self._client.http_request(CONTRACT_CURRENT_URL_2, "get", headers=headers)
@@ -111,7 +135,8 @@ async def fetch_annual_data(self):
         self._logger.info("Fetching annual data")
         await self._client.select_customer(self.account_id, self.customer_id)
         headers = {"Content-Type": "application/json"}
-        res = await self._client.http_request(ANNUAL_DATA_URL, "get", headers=headers)
+        params = {"noContrat": self._selected_contract}
+        res = await self._client.http_request(ANNUAL_DATA_URL, "get", params=params, headers=headers)
         # We can not use res.json() because the response header are not application/json
         json_res = json.loads(await res.text())
         if not json_res.get('results'):
@@ -144,7 +169,8 @@ async def fetch_monthly_data(self):
         self._logger.info("Fetching monthly data")
         await self._client.select_customer(self.account_id, self.customer_id)
         headers = {"Content-Type": "application/json"}
-        res = await self._client.http_request(MONTHLY_DATA_URL, "get", headers=headers)
+        params = {"noContrat": self._selected_contract}
+        res = await self._client.http_request(MONTHLY_DATA_URL, "get", params=params, headers=headers)
         text_res = await res.text()
         # We can not use res.json() because the response header are not application/json
         json_res = json.loads(text_res)
@@ -209,7 +235,10 @@ async def fetch_daily_data(self, start_date=None, end_date=None):
             end_date_str = end_date
 
         headers = {"Content-Type": "application/json"}
-        params = {"dateDebut": start_date_str}
+        params = {
+            "idContrat": self._selected_contract,
+            "dateDebut": start_date_str
+        }
         if end_date_str:
             params.update({"dateFin": end_date_str})
         res = await self._client.http_request(DAILY_DATA_URL, "get",
@@ -266,7 +295,11 @@ async def fetch_hourly_data(self, day=None):
                 return
             day_str = day
 
-        params = {"dateDebut": day_str, "dateFin": day_str}
+        params = {
+            "idContrat": self._selected_contract,
+            "dateDebut": day_str, 
+            "dateFin": day_str
+        }
         res = await self._client.http_request(HOURLY_DATA_URL_2, "get",
                                               params=params, )
         # We can not use res.json() because the response header are not application/json

pyhydroquebec/mqtt_daemon.py

@@ -100,8 +100,9 @@ async def _main_loop(self):
                 # Get contract
                 customer = None
                 for client_customer in client.customers:
-                    if str(client_customer.contract_id) == str(contract_data['id']):
+                    if str(contract_data['id']) in client_customer.contract_list:
                         customer = client_customer
+                        customer.selected_contract = str(contract_data['id'])
 
                 if customer is None:
                     self.logger.warning('Contract %s not found', contract_data['id'])
@@ -132,7 +133,7 @@ async def _main_loop(self):
                 for data_name, data in CURRENT_MAP.items():
                     # Publish sensor
                     sensor_topic = self._publish_sensor(data_name,
-                                                        customer.contract_id,
+                                                        customer.selected_contract,
                                                         unit=data['unit'],
                                                         icon=data['icon'],
                                                         device_class=data['device_class'])
@@ -145,7 +146,7 @@ async def _main_loop(self):
                 for data_name, data in DAILY_MAP.items():
                     # Publish sensor
                     sensor_topic = self._publish_sensor('yesterday_' + data_name,
-                                                        customer.contract_id,
+                                                        customer.selected_contract,
                                                         unit=data['unit'],
                                                         icon=data['icon'],
                                                         device_class=data['device_class'])

tests/test_client.py

@@ -17,7 +17,7 @@ def test_client():
     async_func = client.login()
     loop.run_until_complete(asyncio.gather(async_func))
     assert len(client.customers) > 0
-    assert client.customers[0].contract_id is not None
+    assert client.customers[0].selected_contract is not None
     assert client.customers[0].account_id is not None