From e80964a2bccd4f28a81ef982f6cdfbe49213e6c2 Mon Sep 17 00:00:00 2001
From: Seb Ruiz <seb@sebruiz.net>
Date: Mon, 9 Nov 2020 17:34:02 +1100
Subject: [PATCH] Pass data object into event handler

---
 custom_components/bhyve/__init__.py | 4 ++--
 custom_components/bhyve/sensor.py   | 6 +++---
 custom_components/bhyve/switch.py   | 6 +++---
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/custom_components/bhyve/__init__.py b/custom_components/bhyve/__init__.py
index 60437d1..ff21043 100644
--- a/custom_components/bhyve/__init__.py
+++ b/custom_components/bhyve/__init__.py
@@ -215,7 +215,7 @@ def __init__(
     def _on_ws_data(self, data):
         pass
 
-    def _should_handle_event(self, event_name):
+    def _should_handle_event(self, event_name, data):
         """True if the websocket event should be handled"""
         return True
 
@@ -299,7 +299,7 @@ def update(device_id, data):
                 self._available = False
             elif event == "device_connected":
                 self._available = True
-            if self._should_handle_event(event):
+            if self._should_handle_event(event, data):
                 _LOGGER.info(
                     "Message received: {} - {} - {}".format(
                         self.name, self._device_id, str(data)
diff --git a/custom_components/bhyve/sensor.py b/custom_components/bhyve/sensor.py
index 0e8a993..07a7c86 100644
--- a/custom_components/bhyve/sensor.py
+++ b/custom_components/bhyve/sensor.py
@@ -95,7 +95,7 @@ def unique_id(self):
         """Return a unique, unchanging string that represents this sensor."""
         return f"{self._mac_address}:{self._device_id}:battery"
 
-    def _should_handle_event(self, event_name):
+    def _should_handle_event(self, event_name, data):
         return event_name in [EVENT_CHANGE_MODE]
 
     async def async_update(self):
@@ -139,7 +139,7 @@ def unique_id(self):
         """Return a unique, unchanging string that represents this sensor."""
         return f"{self._mac_address}:{self._device_id}:{self._zone_id}:history"
 
-    def _should_handle_event(self, event_name):
+    def _should_handle_event(self, event_name, data):
         return event_name in [EVENT_DEVICE_IDLE]
 
     async def async_update(self):
@@ -219,5 +219,5 @@ def _on_ws_data(self, data):
         if event == EVENT_CHANGE_MODE:
             self._state = data.get("mode")
 
-    def _should_handle_event(self, event_name):
+    def _should_handle_event(self, event_name, data):
         return event_name in [EVENT_CHANGE_MODE]
diff --git a/custom_components/bhyve/switch.py b/custom_components/bhyve/switch.py
index 2a54598..b20809e 100644
--- a/custom_components/bhyve/switch.py
+++ b/custom_components/bhyve/switch.py
@@ -292,7 +292,7 @@ def _on_ws_data(self, data):
             if program is not None:
                 self._program = program
 
-    def _should_handle_event(self, event_name):
+    def _should_handle_event(self, event_name, data):
         return event_name in [EVENT_PROGRAM_CHANGED]
 
 
@@ -448,7 +448,7 @@ def _set_watering_program(self, program):
                 }
             )
 
-    def _should_handle_event(self, event_name):
+    def _should_handle_event(self, event_name, data):
         return event_name in [
             EVENT_DEVICE_IDLE,
             EVENT_WATERING_COMPLETE,
@@ -576,7 +576,7 @@ def _on_ws_data(self, data):
             # The REST API returns more data about a rain delay (eg cause/weather_type)
             self._update_device_soon()
 
-    def _should_handle_event(self, event_name):
+    def _should_handle_event(self, event_name, data):
         return event_name in [EVENT_RAIN_DELAY]
 
     def _update_device_soon(self):