From 10616f965527ea947e0f04ffe70e7fb07293d6f2 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Tue, 5 Mar 2024 13:26:00 -0500
Subject: [PATCH] Fix config flow for multiple locks (#49)

---
 custom_components/lock_code_manager/__init__.py    |  4 +++-
 custom_components/lock_code_manager/config_flow.py | 13 +++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/custom_components/lock_code_manager/__init__.py b/custom_components/lock_code_manager/__init__.py
index 4a4fe990..020c9141 100644
--- a/custom_components/lock_code_manager/__init__.py
+++ b/custom_components/lock_code_manager/__init__.py
@@ -406,7 +406,9 @@ async def async_update_listener(hass: HomeAssistant, config_entry: ConfigEntry)
                     entry_title,
                     lock,
                 )
-                coordinator = hass_data[entry_id][COORDINATORS][lock_entity_id]
+                coordinator = hass_data[entry_id][COORDINATORS][
+                    lock_entity_id
+                ] = hass_data[COORDINATORS][lock_entity_id]
             else:
                 _LOGGER.debug(
                     "%s (%s): Creating coordinator for lock %s",
diff --git a/custom_components/lock_code_manager/config_flow.py b/custom_components/lock_code_manager/config_flow.py
index 0ae07a14..43f1118c 100644
--- a/custom_components/lock_code_manager/config_flow.py
+++ b/custom_components/lock_code_manager/config_flow.py
@@ -31,6 +31,7 @@
     DEFAULT_START,
     DOMAIN,
 )
+from .data import get_entry_data
 from .helpers import CODE_SLOT_SCHEMA, CODE_SLOTS_SCHEMA, UI_CODE_SLOT_SCHEMA
 
 _LOGGER = logging.getLogger(__name__)
@@ -64,8 +65,12 @@ def _check_common_slots(
             (lock, common_slots, entry.title)
             for lock in locks
             for entry in hass.config_entries.async_entries(DOMAIN)
-            if lock in entry.data[CONF_LOCKS]
-            and (common_slots := sorted(set(entry.data[CONF_SLOTS]) & set(slots_list)))
+            if lock in get_entry_data(entry, CONF_LOCKS, {})
+            and (
+                common_slots := sorted(
+                    set(get_entry_data(entry, CONF_SLOTS, {})) & set(slots_list)
+                )
+            )
             and not (config_entry and config_entry == entry)
         )
     except StopIteration:
@@ -231,7 +236,7 @@ async def async_step_reauth(self, user_input: dict[str, Any]):
             additional_errors, additional_placeholders = _check_common_slots(
                 self.hass,
                 user_input[CONF_LOCKS],
-                config_entry.data[CONF_SLOTS].keys(),
+                get_entry_data(config_entry, CONF_SLOTS, {}).keys(),
                 config_entry,
             )
             errors.update(additional_errors)
@@ -304,7 +309,7 @@ async def async_step_init(
 
         def _get_default(key: str) -> Any:
             """Get default value."""
-            return user_input.get(key, self.config_entry.data[key])
+            return user_input.get(key, get_entry_data(self.config_entry, key, {}))
 
         return self.async_show_form(
             step_id="init",