Pass hass object through to entities

sebr · Jan 19, 2020 · acf547e · acf547e
1 parent 9035224
commit acf547e
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 14 deletions.
diff --git a/custom_components/bhyve/__init__.py b/custom_components/bhyve/__init__.py
@@ -124,9 +124,10 @@ class BHyveEntity(Entity):
     """Define a base BHyve entity."""
 
     def __init__(
-        self, bhyve, device, name, icon, device_class=None,
+        self, hass, bhyve, device, name, icon, device_class=None,
     ):
         """Initialize the sensor."""
+        self._hass = hass
         self._bhyve: Client = bhyve
         self._device_class = device_class
         self._async_unsub_dispatcher_connect = None

diff --git a/custom_components/bhyve/binary_sensor.py b/custom_components/bhyve/binary_sensor.py
@@ -25,7 +25,9 @@ async def async_setup_platform(hass, config, async_add_entities, _discovery_info
                 name = "{0} {1}".format(sensor_type["name"], device.get("name"))
                 _LOGGER.info("Creating binary_sensor: %s", name)
                 binary_sensors.append(
-                    BHyveRainDelayBinarySensor(bhyve, device, name, sensor_type["icon"])
+                    BHyveRainDelayBinarySensor(
+                        hass, bhyve, device, name, sensor_type["icon"]
+                    )
                 )
 
     async_add_entities(binary_sensors, True)
@@ -54,16 +56,16 @@ def _on_ws_data(self, data):
             _LOGGER.warning("No event on ws data {}".format(data))
             return
         elif event == "rain_delay":
-            self._extract_rain_delay(data.get("delay"), {
-                "rain_delay_started_at": data.get("timestamp")
-            })
+            self._extract_rain_delay(
+                data.get("delay"), {"rain_delay_started_at": data.get("timestamp")}
+            )
             # The REST API returns more data about a rain delay (eg cause/weather_type)
             self._update_device_soon()
 
     def _update_device_soon(self):
         if self._update_device_cb is not None:
             self._update_device_cb.cancel()
-        self._update_device_cb = async_call_later(5, self._refetch_device)
+        self._update_device_cb = async_call_later(self._hass, 5, self._refetch_device)
 
     def _extract_rain_delay(self, rain_delay, device_status=None):
         if rain_delay is not None and rain_delay > 0:

diff --git a/custom_components/bhyve/sensor.py b/custom_components/bhyve/sensor.py
@@ -33,6 +33,7 @@ async def async_setup_platform(hass, config, async_add_entities, _discovery_info
                 _LOGGER.info("Creating sensor: %s", name)
                 sensors.append(
                     BHyveSensor(
+                        hass,
                         bhyve,
                         device,
                         name,
@@ -48,9 +49,9 @@ async def async_setup_platform(hass, config, async_add_entities, _discovery_info
 class BHyveSensor(BHyveEntity):
     """Define a BHyve sensor."""
 
-    def __init__(self, bhyve, device, name, icon, unit, device_class):
+    def __init__(self, hass, bhyve, device, name, icon, unit, device_class):
         """Initialize the sensor."""
-        super().__init__(bhyve, device, name, icon, device_class)
+        super().__init__(hass, bhyve, device, name, icon, device_class)
 
         self._unit = unit
 

diff --git a/custom_components/bhyve/switch.py b/custom_components/bhyve/switch.py
@@ -21,6 +21,7 @@
 ATTR_IMAGE_URL = "image_url"
 ATTR_STARTED_WATERING_AT = "started_watering_station_at"
 
+
 async def async_setup_platform(hass, config, async_add_entities, _discovery_info=None):
     """Set up BHyve binary sensors based on a config entry."""
     bhyve = hass.data[DOMAIN]
@@ -35,7 +36,7 @@ async def async_setup_platform(hass, config, async_add_entities, _discovery_info
                 name = "{0} Zone".format(zone.get("name", "Unknown"))
                 _LOGGER.info("Creating switch: %s", name)
                 switches.append(
-                    BHyveSwitch(bhyve, device, zone, name, sensor_type["icon"],)
+                    BHyveSwitch(hass, bhyve, device, zone, name, sensor_type["icon"],)
                 )
 
     async_add_entities(switches, True)
@@ -44,7 +45,7 @@ async def async_setup_platform(hass, config, async_add_entities, _discovery_info
 class BHyveSwitch(BHyveEntity, SwitchDevice):
     """Define a BHyve switch."""
 
-    def __init__(self, bhyve, device, zone, name, icon):
+    def __init__(self, hass, bhyve, device, zone, name, icon):
         """Initialize the switch."""
         self._zone = zone
         self._zone_id = zone.get("station")
@@ -53,7 +54,7 @@ def __init__(self, bhyve, device, zone, name, icon):
             "manual_preset_runtime_sec", DEFAULT_MANUAL_RUNTIME.seconds
         )
 
-        super().__init__(bhyve, device, name, icon, DEVICE_CLASS_SWITCH)
+        super().__init__(hass, bhyve, device, name, icon, DEVICE_CLASS_SWITCH)
 
     def _setup(self, device):
         self._state = None
@@ -79,9 +80,7 @@ def _setup(self, device):
                 and watering_status.get("current_station") == self._zone_id
             )
             self._state = is_watering
-            self._attrs = {
-                ATTR_MANUAL_RUNTIME: self._manual_preset_runtime
-            }
+            self._attrs = {ATTR_MANUAL_RUNTIME: self._manual_preset_runtime}
 
             smart_watering_enabled = zone.get("smart_watering_enabled")
             if smart_watering_enabled is not None: