Fixed having multiple contracts with zps

custom_components/wnsm/api/client.py

@@ -238,15 +238,19 @@ def _call_api(
         return response.json()
 
     def get_zaehlpunkt(self, zaehlpunkt: str = None) -> tuple[str, str, str]:
-        zps = self.zaehlpunkte()[0]
-        customer_id = zps["geschaeftspartner"]
-        if zaehlpunkt is not None:
-            zp = [z for z in zps["zaehlpunkte"] if z["zaehlpunktnummer"] == zaehlpunkt]
-            anlagetype = zp[0]["anlage"]["typ"] if len(zp) > 0 else None
-            zp = zp[0]["zaehlpunktnummer"] if len(zp) > 0 else None
+        contracts = self.zaehlpunkte()
+        if zaehlpunkt is None:
+            customer_id = contracts[0]["geschaeftspartner"]
+            zp = contracts[0]["zaehlpunkte"][0]["zaehlpunktnummer"]
+            anlagetype = contracts[0]["zaehlpunkte"][0]["anlage"]["typ"]
         else:
-            zp = zps["zaehlpunkte"][0]["zaehlpunktnummer"]
-            anlagetype = zps["zaehlpunkte"][0]["anlage"]["typ"]
+            customer_id = zp = anlagetype = None
+            for contract in contracts:
+                zp = [z for z in contract["zaehlpunkte"] if z["zaehlpunktnummer"] == zaehlpunkt]
+                if len(zp) > 0:
+                    anlagetype = zp[0]["anlage"]["typ"]
+                    zp = zp[0]["zaehlpunktnummer"]
+                    customer_id = contract["geschaeftspartner"]
         return customer_id, zp, const.AnlageType.from_str(anlagetype)
 
     def zaehlpunkte(self):

custom_components/wnsm/base_sensor.py

@@ -86,16 +86,17 @@ def available(self) -> bool:
     def state(self) -> Optional[str]:  # pylint: disable=overridden-final-method
         return self._state
 
-    def contracts2zaehlpunkte(self, contracts: dict) -> [dict]:
-        if contracts is None or len(contracts) == 0:
-            raise RuntimeError(f"Cannot access Zaehlpunkt {self.zaehlpunkt}")
-        geschaeftspartner = contracts[0]["geschaeftspartner"] if "geschaeftspartner" in contracts[0] else None
-        if contracts is not None and isinstance(contracts, list) and len(contracts) > 0 and "zaehlpunkte" in contracts[0]:
-            zaehlpunkte = [
-                {**z, "geschaeftspartner": geschaeftspartner} for z in contracts[0]["zaehlpunkte"] if z["zaehlpunktnummer"] == self.zaehlpunkt
-            ]
+    def contracts2zaehlpunkte(self, contracts: dict) -> list[dict]:
+        zaehlpunkte = []
+        if contracts is not None and isinstance(contracts, list) and len(contracts) > 0:
+            for contract in contracts:
+                if "zaehlpunkte" in contract:
+                    geschaeftspartner = contract["geschaeftspartner"] if "geschaeftspartner" in contract else None
+                    zaehlpunkte += [
+                        {**z, "geschaeftspartner": geschaeftspartner} for z in contract["zaehlpunkte"] if z["zaehlpunktnummer"] == self.zaehlpunkt
+                    ]
         else:
-            zaehlpunkte = []
+            raise RuntimeError(f"Cannot access Zaehlpunkt {self.zaehlpunkt}")
         return zaehlpunkte
 
     async def get_zaehlpunkt(self, smartmeter: Smartmeter) -> dict[str, str]:

custom_components/wnsm/config_flow.py

@@ -32,10 +32,14 @@ async def validate_auth(self, username: str, password: str) -> list[dict]:
         """
         smartmeter = Smartmeter(username, password)
         await self.hass.async_add_executor_job(smartmeter.login)
-        zps = await self.hass.async_add_executor_job(smartmeter.zaehlpunkte)
-        if zps is not None and isinstance(zps, list) and len(zps) > 0 and "zaehlpunkte" in zps[0]:
-            return zps[0]["zaehlpunkte"]
-        return []
+        contracts = await self.hass.async_add_executor_job(smartmeter.zaehlpunkte)
+        zaehlpunkte=[]
+        if contracts is not None and isinstance(contracts, list) and len(contracts) > 0:
+            for contract in contracts:
+                if "zaehlpunkte" in contract:
+                    zaehlpunkte+=contract["zaehlpunkte"]
+        return zaehlpunkte
+
 
     async def async_step_user(self, user_input: Optional[dict[str, Any]] = None):
         """Invoked when a user initiates a flow via the user interface."""