From 93b523141894152bb4a97d791f49ecda2a044283 Mon Sep 17 00:00:00 2001
From: lunDreame <87955512+lunDreame@users.noreply.github.com>
Date: Thu, 2 Jan 2025 00:58:01 +0900
Subject: [PATCH] Patch Update

---
 custom_components/kocom_wallpad/__init__.py | 5 ++++-
 custom_components/kocom_wallpad/entity.py   | 9 +++++----
 custom_components/kocom_wallpad/gateway.py  | 8 ++++++--
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/custom_components/kocom_wallpad/__init__.py b/custom_components/kocom_wallpad/__init__.py
index ce1df20..f93da03 100644
--- a/custom_components/kocom_wallpad/__init__.py
+++ b/custom_components/kocom_wallpad/__init__.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from homeassistant.const import Platform
+from homeassistant.const import Platform, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
 from homeassistant.config_entries import ConfigEntry
 
@@ -32,6 +32,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     await gateway.async_start()
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+    #entry.async_on_unload(
+    #    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, gateway.async_close)
+    #)
 
     return True
 
diff --git a/custom_components/kocom_wallpad/entity.py b/custom_components/kocom_wallpad/entity.py
index d6ff390..c71712a 100644
--- a/custom_components/kocom_wallpad/entity.py
+++ b/custom_components/kocom_wallpad/entity.py
@@ -34,7 +34,7 @@ class KocomEntity(RestoreEntity):
     """Base class for Kocom Wallpad entities."""
 
     _attr_has_entity_name = True
-    _attr_should_poll = True
+    _attr_should_poll = False
     
     def __init__(
         self,
@@ -85,10 +85,11 @@ def available(self) -> bool:
         return self.gateway.connection.is_connected()
     
     @callback
-    def async_handle_device_update(self, device: Device) -> None:
+    def async_handle_device_update(self, packet: KocomPacket) -> None:
         """Handle device update."""
-        if self.device.state != device.state:
-            self.device.state = device.state
+        if self.packet._device.state != packet._device.state:
+            self.packet = packet
+            self.device = packet._device
             self.async_write_ha_state()
 
     async def async_added_to_hass(self) -> None:
diff --git a/custom_components/kocom_wallpad/gateway.py b/custom_components/kocom_wallpad/gateway.py
index 41a5136..af37c9d 100644
--- a/custom_components/kocom_wallpad/gateway.py
+++ b/custom_components/kocom_wallpad/gateway.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from homeassistant.const import Platform, CONF_HOST, CONF_PORT
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, Event
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers import entity_registry as er, restore_state
@@ -60,6 +60,10 @@ async def async_start(self) -> None:
         await self.client.start()
         self.client.add_device_callback(self._handle_device_update)
         
+    async def async_close(self, event: Event) -> None:
+        """Close the gateway."""
+        await self.async_disconnect()
+    
     def get_entities(self, platform: Platform) -> list[KocomPacket]:
         """Get the entities for the platform."""
         return list(self.entities.get(platform, {}).values())
@@ -105,7 +109,7 @@ async def _handle_device_update(self, packet: KocomPacket) -> None:
             async_dispatcher_send(self.hass, add_signal, packet)
         
         device_update_signal = f"{DOMAIN}_{self.host}_{dev_id}"
-        async_dispatcher_send(self.hass, device_update_signal, device)
+        async_dispatcher_send(self.hass, device_update_signal, packet)
         
     def parse_platform(self, packet: KocomPacket) -> Platform | None:
         """Parse the platform from the packet."""