From d175e19c3bc85243641cccab737831b53104f4ae Mon Sep 17 00:00:00 2001
From: Andre Basche <andre.basche@gmail.com>
Date: Sun, 11 Feb 2024 05:06:53 +0100
Subject: [PATCH] Safe refresh token

---
 custom_components/hon/__init__.py | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/custom_components/hon/__init__.py b/custom_components/hon/__init__.py
index 24b97a3..c4f9aea 100644
--- a/custom_components/hon/__init__.py
+++ b/custom_components/hon/__init__.py
@@ -29,16 +29,22 @@ async def async_setup_entry(hass: HomeAssistantType, entry: ConfigEntry) -> bool
     session = aiohttp_client.async_get_clientsession(hass)
     if (config_dir := hass.config.config_dir) is None:
         raise ValueError("Missing Config Dir")
-    hon = await Hon(
-        entry.data[CONF_EMAIL],
-        entry.data[CONF_PASSWORD],
-        mobile_id=MOBILE_ID,
-        session=session,
-        refresh_token=entry.data[CONF_REFRESH_TOKEN],
-        test_data_path=Path(config_dir),
-    ).create()
+    kwargs = {
+        "email": entry.data[CONF_EMAIL],
+        "password": entry.data[CONF_PASSWORD],
+        "mobile_id": MOBILE_ID,
+        "session": session,
+        "test_data_path": Path(config_dir),
+    }
+    if refresh_token := entry.data.get(CONF_REFRESH_TOKEN):
+        kwargs["refresh_token"] = refresh_token
+    hon = await Hon(**kwargs).create()
     hass.data.setdefault(DOMAIN, {})
     hass.data[DOMAIN][entry.unique_id] = hon
+
+    hass.config_entries.async_update_entry(
+        entry, data={**entry.data, CONF_REFRESH_TOKEN: hon.api.auth.refresh_token}
+    )
     hass.data[DOMAIN]["coordinators"] = {}
 
     for platform in PLATFORMS:
@@ -49,7 +55,11 @@ async def async_setup_entry(hass: HomeAssistantType, entry: ConfigEntry) -> bool
 
 
 async def async_unload_entry(hass: HomeAssistantType, entry: ConfigEntry) -> bool:
-    entry.data[CONF_REFRESH_TOKEN] = hass.data[DOMAIN][entry.unique_id].api.auth.refresh_token
+    refresh_token = hass.data[DOMAIN][entry.unique_id].api.auth.refresh_token
+
+    hass.config_entries.async_update_entry(
+        entry, data={**entry.data, CONF_REFRESH_TOKEN: refresh_token}
+    )
     unload = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload:
         if not hass.data[DOMAIN]: