From 8e75547ca461f66135909cfb016aa682121cd15a Mon Sep 17 00:00:00 2001
From: Aaron Bach <bachya1208@gmail.com>
Date: Mon, 30 May 2022 15:36:58 -0600
Subject: [PATCH] Guard against missing data in 1st generation RainMachine
 controllers (#72632)

---
 .../components/rainmachine/binary_sensor.py   | 16 +++----
 .../components/rainmachine/sensor.py          |  2 +-
 .../components/rainmachine/switch.py          | 43 +++++++++++--------
 3 files changed, 35 insertions(+), 26 deletions(-)

diff --git a/homeassistant/components/rainmachine/binary_sensor.py b/homeassistant/components/rainmachine/binary_sensor.py
index fb404adb199..730b51c142a 100644
--- a/homeassistant/components/rainmachine/binary_sensor.py
+++ b/homeassistant/components/rainmachine/binary_sensor.py
@@ -158,17 +158,17 @@ class CurrentRestrictionsBinarySensor(RainMachineEntity, BinarySensorEntity):
     def update_from_latest_data(self) -> None:
         """Update the state."""
         if self.entity_description.key == TYPE_FREEZE:
-            self._attr_is_on = self.coordinator.data["freeze"]
+            self._attr_is_on = self.coordinator.data.get("freeze")
         elif self.entity_description.key == TYPE_HOURLY:
-            self._attr_is_on = self.coordinator.data["hourly"]
+            self._attr_is_on = self.coordinator.data.get("hourly")
         elif self.entity_description.key == TYPE_MONTH:
-            self._attr_is_on = self.coordinator.data["month"]
+            self._attr_is_on = self.coordinator.data.get("month")
         elif self.entity_description.key == TYPE_RAINDELAY:
-            self._attr_is_on = self.coordinator.data["rainDelay"]
+            self._attr_is_on = self.coordinator.data.get("rainDelay")
         elif self.entity_description.key == TYPE_RAINSENSOR:
-            self._attr_is_on = self.coordinator.data["rainSensor"]
+            self._attr_is_on = self.coordinator.data.get("rainSensor")
         elif self.entity_description.key == TYPE_WEEKDAY:
-            self._attr_is_on = self.coordinator.data["weekDay"]
+            self._attr_is_on = self.coordinator.data.get("weekDay")
 
 
 class ProvisionSettingsBinarySensor(RainMachineEntity, BinarySensorEntity):
@@ -188,6 +188,6 @@ class UniversalRestrictionsBinarySensor(RainMachineEntity, BinarySensorEntity):
     def update_from_latest_data(self) -> None:
         """Update the state."""
         if self.entity_description.key == TYPE_FREEZE_PROTECTION:
-            self._attr_is_on = self.coordinator.data["freezeProtectEnabled"]
+            self._attr_is_on = self.coordinator.data.get("freezeProtectEnabled")
         elif self.entity_description.key == TYPE_HOT_DAYS:
-            self._attr_is_on = self.coordinator.data["hotDaysExtraWatering"]
+            self._attr_is_on = self.coordinator.data.get("hotDaysExtraWatering")
diff --git a/homeassistant/components/rainmachine/sensor.py b/homeassistant/components/rainmachine/sensor.py
index b825faca7e1..a2b0f7cd539 100644
--- a/homeassistant/components/rainmachine/sensor.py
+++ b/homeassistant/components/rainmachine/sensor.py
@@ -198,7 +198,7 @@ class UniversalRestrictionsSensor(RainMachineEntity, SensorEntity):
     def update_from_latest_data(self) -> None:
         """Update the state."""
         if self.entity_description.key == TYPE_FREEZE_TEMP:
-            self._attr_native_value = self.coordinator.data["freezeProtectTemp"]
+            self._attr_native_value = self.coordinator.data.get("freezeProtectTemp")
 
 
 class ZoneTimeRemainingSensor(RainMachineEntity, SensorEntity):
diff --git a/homeassistant/components/rainmachine/switch.py b/homeassistant/components/rainmachine/switch.py
index 007aec97a3e..a220aafa2a5 100644
--- a/homeassistant/components/rainmachine/switch.py
+++ b/homeassistant/components/rainmachine/switch.py
@@ -389,23 +389,32 @@ class RainMachineZone(RainMachineActivitySwitch):
 
         self._attr_is_on = bool(data["state"])
 
-        self._attr_extra_state_attributes.update(
-            {
-                ATTR_AREA: round(data["waterSense"]["area"], 2),
-                ATTR_CURRENT_CYCLE: data["cycle"],
-                ATTR_FIELD_CAPACITY: round(data["waterSense"]["fieldCapacity"], 2),
-                ATTR_ID: data["uid"],
-                ATTR_NO_CYCLES: data["noOfCycles"],
-                ATTR_PRECIP_RATE: round(data["waterSense"]["precipitationRate"], 2),
-                ATTR_RESTRICTIONS: data["restriction"],
-                ATTR_SLOPE: SLOPE_TYPE_MAP.get(data["slope"], 99),
-                ATTR_SOIL_TYPE: SOIL_TYPE_MAP.get(data["soil"], 99),
-                ATTR_SPRINKLER_TYPE: SPRINKLER_TYPE_MAP.get(data["group_id"], 99),
-                ATTR_STATUS: RUN_STATE_MAP[data["state"]],
-                ATTR_SUN_EXPOSURE: SUN_EXPOSURE_MAP.get(data.get("sun")),
-                ATTR_VEGETATION_TYPE: VEGETATION_MAP.get(data["type"], 99),
-            }
-        )
+        attrs = {
+            ATTR_CURRENT_CYCLE: data["cycle"],
+            ATTR_ID: data["uid"],
+            ATTR_NO_CYCLES: data["noOfCycles"],
+            ATTR_RESTRICTIONS: data("restriction"),
+            ATTR_SLOPE: SLOPE_TYPE_MAP.get(data["slope"], 99),
+            ATTR_SOIL_TYPE: SOIL_TYPE_MAP.get(data["soil"], 99),
+            ATTR_SPRINKLER_TYPE: SPRINKLER_TYPE_MAP.get(data["group_id"], 99),
+            ATTR_STATUS: RUN_STATE_MAP[data["state"]],
+            ATTR_SUN_EXPOSURE: SUN_EXPOSURE_MAP.get(data.get("sun")),
+            ATTR_VEGETATION_TYPE: VEGETATION_MAP.get(data["type"], 99),
+        }
+
+        if "waterSense" in data:
+            if "area" in data["waterSense"]:
+                attrs[ATTR_AREA] = round(data["waterSense"]["area"], 2)
+            if "fieldCapacity" in data["waterSense"]:
+                attrs[ATTR_FIELD_CAPACITY] = round(
+                    data["waterSense"]["fieldCapacity"], 2
+                )
+            if "precipitationRate" in data["waterSense"]:
+                attrs[ATTR_PRECIP_RATE] = round(
+                    data["waterSense"]["precipitationRate"], 2
+                )
+
+        self._attr_extra_state_attributes.update(attrs)
 
 
 class RainMachineZoneEnabled(RainMachineEnabledSwitch):