From a40ddb5e838d62095ea1dc1c60fec9a41cf46491 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Wed, 24 Aug 2022 09:11:01 +0200
Subject: [PATCH] Use _attr_should_poll in xiaomi_aqara entities (#77197)

* Use _attr_should_poll in xiaomi_aqara entities

* Adjust switch
---
 homeassistant/components/xiaomi_aqara/__init__.py  |  7 ++-----
 .../components/xiaomi_aqara/binary_sensor.py       | 14 ++++----------
 homeassistant/components/xiaomi_aqara/switch.py    |  7 ++-----
 3 files changed, 8 insertions(+), 20 deletions(-)

diff --git a/homeassistant/components/xiaomi_aqara/__init__.py b/homeassistant/components/xiaomi_aqara/__init__.py
index da3303494e5b..e99851dae9f8 100644
--- a/homeassistant/components/xiaomi_aqara/__init__.py
+++ b/homeassistant/components/xiaomi_aqara/__init__.py
@@ -225,6 +225,8 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 class XiaomiDevice(Entity):
     """Representation a base Xiaomi device."""
 
+    _attr_should_poll = False
+
     def __init__(self, device, device_type, xiaomi_hub, config_entry):
         """Initialize the Xiaomi device."""
         self._state = None
@@ -309,11 +311,6 @@ class XiaomiDevice(Entity):
         """Return True if entity is available."""
         return self._is_available
 
-    @property
-    def should_poll(self):
-        """Return the polling state. No polling needed."""
-        return False
-
     @property
     def extra_state_attributes(self):
         """Return the state attributes."""
diff --git a/homeassistant/components/xiaomi_aqara/binary_sensor.py b/homeassistant/components/xiaomi_aqara/binary_sensor.py
index 04e3945e7a2c..44c17b634cb2 100644
--- a/homeassistant/components/xiaomi_aqara/binary_sensor.py
+++ b/homeassistant/components/xiaomi_aqara/binary_sensor.py
@@ -140,15 +140,9 @@ class XiaomiBinarySensor(XiaomiDevice, BinarySensorEntity):
         """Initialize the XiaomiSmokeSensor."""
         self._data_key = data_key
         self._device_class = device_class
-        self._should_poll = False
         self._density = 0
         super().__init__(device, name, xiaomi_hub, config_entry)
 
-    @property
-    def should_poll(self):
-        """Return True if entity has to be polled for state."""
-        return self._should_poll
-
     @property
     def is_on(self):
         """Return true if sensor is on."""
@@ -340,7 +334,7 @@ class XiaomiDoorSensor(XiaomiBinarySensor, RestoreEntity):
 
     def parse_data(self, data, raw_data):
         """Parse data sent by gateway."""
-        self._should_poll = False
+        self._attr_should_poll = False
         if NO_CLOSE in data:  # handle push from the hub
             self._open_since = data[NO_CLOSE]
             return True
@@ -350,7 +344,7 @@ class XiaomiDoorSensor(XiaomiBinarySensor, RestoreEntity):
             return False
 
         if value == "open":
-            self._should_poll = True
+            self._attr_should_poll = True
             if self._state:
                 return False
             self._state = True
@@ -388,14 +382,14 @@ class XiaomiWaterLeakSensor(XiaomiBinarySensor):
 
     def parse_data(self, data, raw_data):
         """Parse data sent by gateway."""
-        self._should_poll = False
+        self._attr_should_poll = False
 
         value = data.get(self._data_key)
         if value is None:
             return False
 
         if value == "leak":
-            self._should_poll = True
+            self._attr_should_poll = True
             if self._state:
                 return False
             self._state = True
diff --git a/homeassistant/components/xiaomi_aqara/switch.py b/homeassistant/components/xiaomi_aqara/switch.py
index 86acd4100a2a..39bf637256f4 100644
--- a/homeassistant/components/xiaomi_aqara/switch.py
+++ b/homeassistant/components/xiaomi_aqara/switch.py
@@ -149,6 +149,8 @@ class XiaomiGenericSwitch(XiaomiDevice, SwitchEntity):
         self._load_power = None
         self._power_consumed = None
         self._supports_power_consumption = supports_power_consumption
+        # Polling needed for Zigbee plug only.
+        self._attr_should_poll = supports_power_consumption
         super().__init__(device, name, xiaomi_hub, config_entry)
 
     @property
@@ -177,11 +179,6 @@ class XiaomiGenericSwitch(XiaomiDevice, SwitchEntity):
         attrs.update(super().extra_state_attributes)
         return attrs
 
-    @property
-    def should_poll(self):
-        """Return the polling state. Polling needed for Zigbee plug only."""
-        return self._supports_power_consumption
-
     def turn_on(self, **kwargs):
         """Turn the switch on."""
         if self._write_to_hub(self._sid, **{self._data_key: "on"}):