From be85e29aa0abf48b29020a8323f8280fda432d72 Mon Sep 17 00:00:00 2001
From: Michal4K <michal.jal@outlook.com>
Date: Mon, 2 Sep 2024 22:34:24 +0000
Subject: [PATCH] use senor_type instead of switchbot model + tests

---
 .../components/switchbot/config_flow.py       |  4 +-
 .../components/switchbot/test_config_flow.py  | 63 +++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/switchbot/config_flow.py b/homeassistant/components/switchbot/config_flow.py
index a8b894cc7c920..0468db5618adf 100644
--- a/homeassistant/components/switchbot/config_flow.py
+++ b/homeassistant/components/switchbot/config_flow.py
@@ -11,7 +11,6 @@
     SwitchbotApiError,
     SwitchbotAuthenticationError,
     SwitchbotLock,
-    SwitchbotModel,
     parse_advertisement_data,
 )
 import voluptuous as vol
@@ -48,6 +47,7 @@
     NON_CONNECTABLE_SUPPORTED_MODEL_TYPES,
     SUPPORTED_LOCK_MODELS,
     SUPPORTED_MODEL_TYPES,
+    SupportedModels,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -366,7 +366,7 @@ async def async_step_init(
                 ),
             ): int
         }
-        if self.config_entry.runtime_data.model == SwitchbotModel.LOCK_PRO:
+        if self.config_entry.data.get(CONF_SENSOR_TYPE) == SupportedModels.LOCK_PRO:
             options.update(
                 {
                     vol.Optional(
diff --git a/tests/components/switchbot/test_config_flow.py b/tests/components/switchbot/test_config_flow.py
index 182e9457f223b..b0fba2a5f18ab 100644
--- a/tests/components/switchbot/test_config_flow.py
+++ b/tests/components/switchbot/test_config_flow.py
@@ -7,6 +7,7 @@
 from homeassistant.components.switchbot.const import (
     CONF_ENCRYPTION_KEY,
     CONF_KEY_ID,
+    CONF_LOCK_NIGHTLATCH,
     CONF_RETRY_COUNT,
 )
 from homeassistant.config_entries import SOURCE_BLUETOOTH, SOURCE_USER
@@ -782,3 +783,65 @@ async def test_options_flow(hass: HomeAssistant) -> None:
     assert len(mock_setup_entry.mock_calls) == 1
 
     assert entry.options[CONF_RETRY_COUNT] == 6
+
+
+async def test_options_flow_lock_pro(hass: HomeAssistant) -> None:
+    """Test updating options."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_ADDRESS: "aa:bb:cc:dd:ee:ff",
+            CONF_NAME: "test-name",
+            CONF_PASSWORD: "test-password",
+            CONF_SENSOR_TYPE: "lock_pro",
+        },
+        options={CONF_RETRY_COUNT: 10},
+        unique_id="aabbccddeeff",
+    )
+    entry.add_to_hass(hass)
+
+    # Test Force night_latch should be disabled by default.
+    with patch_async_setup_entry() as mock_setup_entry:
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+        result = await hass.config_entries.options.async_init(entry.entry_id)
+        assert result["type"] is FlowResultType.FORM
+        assert result["step_id"] == "init"
+        assert result["errors"] is None
+
+        result = await hass.config_entries.options.async_configure(
+            result["flow_id"],
+            user_input={
+                CONF_RETRY_COUNT: 3,
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["data"][CONF_LOCK_NIGHTLATCH] is False
+
+    assert len(mock_setup_entry.mock_calls) == 1
+
+    # Test Set force night_latch to be enabled.
+
+    with patch_async_setup_entry() as mock_setup_entry:
+        result = await hass.config_entries.options.async_init(entry.entry_id)
+        assert result["type"] is FlowResultType.FORM
+        assert result["step_id"] == "init"
+        assert result["errors"] is None
+
+        result = await hass.config_entries.options.async_configure(
+            result["flow_id"],
+            user_input={
+                CONF_LOCK_NIGHTLATCH: True,
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["data"][CONF_LOCK_NIGHTLATCH] is True
+
+    assert len(mock_setup_entry.mock_calls) == 0
+
+    assert entry.options[CONF_LOCK_NIGHTLATCH] is True