fix: Added more robust checks when retrieving account information

BottlecapDave · Apr 1, 2023 · 5898f41 · 5898f41
1 parent e586213
commit 5898f41
Showing 1 changed file with 35 additions and 27 deletions.
diff --git a/custom_components/octopus_energy/api_client.py b/custom_components/octopus_energy/api_client.py
@@ -174,34 +174,42 @@ async def async_get_account(self, account_id):
             account_response_body["data"]["account"] is not None):
           return {
             "electricity_meter_points": list(map(lambda mp: {
-              "mpan": mp["meterPoint"]["mpan"],
-              "meters": list(map(lambda m: {
-                "serial_number": m["serialNumber"],
-                "is_export": m["smartExportElectricityMeter"] != None,
-                "is_smart_meter": m["smartImportElectricityMeter"] != None or m["smartExportElectricityMeter"] != None,
-                "device_id": m["smartImportElectricityMeter"]["deviceId"] if m["smartImportElectricityMeter"] != None else None
-              }, mp["meterPoint"]["meters"])),
-              "agreements": list(map(lambda a: {
-                "valid_from": a["validFrom"],
-                "valid_to": a["validTo"],
-                "tariff_code": a["tariff"]["tariffCode"] if "tariff" in a and "tariffCode" in a["tariff"] else None,
-                "product_code": a["tariff"]["productCode"] if "tariff" in a and "productCode" in a["tariff"] else None,
-              }, mp["meterPoint"]["agreements"]))
-            }, account_response_body["data"]["account"]["electricityAgreements"])),
+                "mpan": mp["meterPoint"]["mpan"],
+                "meters": list(map(lambda m: {
+                  "serial_number": m["serialNumber"],
+                  "is_export": m["smartExportElectricityMeter"] != None,
+                  "is_smart_meter": m["smartImportElectricityMeter"] != None or m["smartExportElectricityMeter"] != None,
+                  "device_id": m["smartImportElectricityMeter"]["deviceId"] if m["smartImportElectricityMeter"] != None else None
+                }, mp["meterPoint"]["meters"])),
+                "agreements": list(map(lambda a: {
+                  "valid_from": a["validFrom"],
+                  "valid_to": a["validTo"],
+                  "tariff_code": a["tariff"]["tariffCode"] if "tariff" in a and "tariffCode" in a["tariff"] else None,
+                  "product_code": a["tariff"]["productCode"] if "tariff" in a and "productCode" in a["tariff"] else None,
+                }, mp["meterPoint"]["agreements"]))
+              }, 
+              account_response_body["data"]["account"]["electricityAgreements"]
+              if "electricityAgreements" in account_response_body["data"]["account"] and account_response_body["data"]["account"]["electricityAgreements"] is not None
+              else []
+            )),
             "gas_meter_points": list(map(lambda mp: {
-              "mprn": mp["meterPoint"]["mprn"],
-              "meters": list(map(lambda m: {
-                "serial_number": m["serialNumber"],
-                "consumption_units": m["consumptionUnits"],
-                "device_id": m["smartGasMeter"]["deviceId"] if m["smartGasMeter"] != None else None
-              }, mp["meterPoint"]["meters"])),
-              "agreements": list(map(lambda a: {
-                "valid_from": a["validFrom"],
-                "valid_to": a["validTo"],
-                "tariff_code": a["tariff"]["tariffCode"] if "tariff" in a and "tariffCode" in a["tariff"] else None,
-                "product_code": a["tariff"]["productCode"] if "tariff" in a and "productCode" in a["tariff"] else None,
-              }, mp["meterPoint"]["agreements"]))
-            }, account_response_body["data"]["account"]["gasAgreements"])),
+                "mprn": mp["meterPoint"]["mprn"],
+                "meters": list(map(lambda m: {
+                  "serial_number": m["serialNumber"],
+                  "consumption_units": m["consumptionUnits"],
+                  "device_id": m["smartGasMeter"]["deviceId"] if m["smartGasMeter"] != None else None
+                }, mp["meterPoint"]["meters"])),
+                "agreements": list(map(lambda a: {
+                  "valid_from": a["validFrom"],
+                  "valid_to": a["validTo"],
+                  "tariff_code": a["tariff"]["tariffCode"] if "tariff" in a and "tariffCode" in a["tariff"] else None,
+                  "product_code": a["tariff"]["productCode"] if "tariff" in a and "productCode" in a["tariff"] else None,
+                }, mp["meterPoint"]["agreements"]))
+              }, 
+              account_response_body["data"]["account"]["gasAgreements"] 
+              if "gasAgreements" in account_response_body["data"]["account"] and account_response_body["data"]["account"]["gasAgreements"] is not None
+              else []
+            )),
           }
         else:
           _LOGGER.error("Failed to retrieve account")