From c9245376926ff19efceb34637e6372e3980fbe9d Mon Sep 17 00:00:00 2001
From: romedtino <romedtino@gmail.com>
Date: Mon, 24 Jul 2023 22:24:21 -0700
Subject: [PATCH] Fix for refresh token

---
 custom_components/simple_wyze_vac/const.py      |  3 ++-
 custom_components/simple_wyze_vac/manifest.json |  2 +-
 custom_components/simple_wyze_vac/vacuum.py     | 12 ++++++++----
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/custom_components/simple_wyze_vac/const.py b/custom_components/simple_wyze_vac/const.py
index 12d6815..574283b 100644
--- a/custom_components/simple_wyze_vac/const.py
+++ b/custom_components/simple_wyze_vac/const.py
@@ -6,4 +6,5 @@
 WYZE_SCAN_INTERVAL = "simple_wyze_vac_scan_interval"
 CONF_POLLING = "poll"
 CONF_TOTP = "totp"
-CONF_KEY_ID = "key_id"
\ No newline at end of file
+CONF_KEY_ID = "key_id"
+CONF_API_KEY = "api_key"
\ No newline at end of file
diff --git a/custom_components/simple_wyze_vac/manifest.json b/custom_components/simple_wyze_vac/manifest.json
index 653dd5c..b6c2997 100644
--- a/custom_components/simple_wyze_vac/manifest.json
+++ b/custom_components/simple_wyze_vac/manifest.json
@@ -7,5 +7,5 @@
   "iot_class": "cloud_polling",
   "name": "Simple Wyze Vacuum",
   "requirements": ["https://github.com/romedtino/wyze-sdk/archive/refs/tags/10.4.1.zip#wyze-sdk==10.4.1"],
-  "version": "1.8.8"
+  "version": "1.8.9"
 }
diff --git a/custom_components/simple_wyze_vac/vacuum.py b/custom_components/simple_wyze_vac/vacuum.py
index 512411c..4f0a1b8 100644
--- a/custom_components/simple_wyze_vac/vacuum.py
+++ b/custom_components/simple_wyze_vac/vacuum.py
@@ -10,7 +10,7 @@
 
 from .const import CONF_TOTP, WYZE_VAC_CLIENT, WYZE_VACUUMS, WYZE_USERNAME, WYZE_PASSWORD, \
                    DOMAIN, \
-                   CONF_POLLING, WYZE_SCAN_INTERVAL
+                   CONF_POLLING, WYZE_SCAN_INTERVAL, CONF_KEY_ID, CONF_API_KEY
 
 from wyze_sdk.models.devices import VacuumMode, VacuumSuctionLevel
 from wyze_sdk.errors import WyzeApiError, WyzeClientNotConnectedError
@@ -66,6 +66,8 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 
     username = hass.data[WYZE_USERNAME]
     password = hass.data[WYZE_PASSWORD]
+    key_id = hass.data[CONF_KEY_ID]
+    api_key = hass.data[CONF_API_KEY]
     polling = hass.data[CONF_POLLING]
     totp = hass.data[CONF_TOTP]
 
@@ -81,7 +83,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 
     vacs = []
     for pl in hass.data[WYZE_VACUUMS]:
-        vacs.append(WyzeVac(client, pl, username, password, totp, polling, scan_interval))
+        vacs.append(WyzeVac(client, pl, username, password, key_id, api_key, totp, polling, scan_interval))
 
     def refresh(event_time):
         """Refresh"""
@@ -108,7 +110,7 @@ def refresh(event_time):
 
 class WyzeVac(StateVacuumEntity):
 
-    def __init__(self, client, pl, username, password, totp, polling, scan_interval):
+    def __init__(self, client, pl, username, password, key_id, api_key, totp, polling, scan_interval):
         self._client = client
         self._vac_mac = pl["mac"]
         self._model = pl["model"]
@@ -123,6 +125,8 @@ def __init__(self, client, pl, username, password, totp, polling, scan_interval)
 
         self._username = username
         self._password = password
+        self._key_id = key_id
+        self._api_key = api_key
         self._totp = totp
 
         self._polling = polling
@@ -217,7 +221,7 @@ def extra_state_attributes(self):
 
     async def get_new_client(self):
         _LOGGER.warn("Refreshing Wyze Client. Do this sparingly to be prevent lockout.")
-        self._client = await self.hass.async_add_executor_job(lambda: Client(email=self._username, password=self._password, totp_key=self._totp))
+        self._client = await self.hass.async_add_executor_job(lambda: Client(email=self._username, password=self._password, key_id=self._key_id, api_key=self._api_key, totp_key=self._totp))
 
     async def async_start(self, **kwargs):
         try: