From f7909ee34a1514e7dc423a19f543f9fd6443681f Mon Sep 17 00:00:00 2001
From: Simone Chemelli <simone.chemelli@gmail.com>
Date: Tue, 30 Jan 2024 13:19:40 +0100
Subject: [PATCH] Clean up Fritz options flow (#109111)

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>
Co-authored-by: Franck Nijhof <frenck@frenck.nl>
---
 homeassistant/components/fritz/config_flow.py | 19 +++++-----
 tests/components/fritz/test_config_flow.py    | 36 +++++++++----------
 2 files changed, 27 insertions(+), 28 deletions(-)

diff --git a/homeassistant/components/fritz/config_flow.py b/homeassistant/components/fritz/config_flow.py
index 5474a171798..03bcc3b77f7 100644
--- a/homeassistant/components/fritz/config_flow.py
+++ b/homeassistant/components/fritz/config_flow.py
@@ -17,7 +17,12 @@ from homeassistant.components.device_tracker import (
     CONF_CONSIDER_HOME,
     DEFAULT_CONSIDER_HOME,
 )
-from homeassistant.config_entries import ConfigEntry, ConfigFlow, OptionsFlow
+from homeassistant.config_entries import (
+    ConfigEntry,
+    ConfigFlow,
+    OptionsFlow,
+    OptionsFlowWithConfigEntry,
+)
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
 from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
@@ -289,12 +294,8 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
         return self.async_abort(reason="reauth_successful")
 
 
-class FritzBoxToolsOptionsFlowHandler(OptionsFlow):
-    """Handle a option flow."""
-
-    def __init__(self, config_entry: ConfigEntry) -> None:
-        """Initialize options flow."""
-        self.config_entry = config_entry
+class FritzBoxToolsOptionsFlowHandler(OptionsFlowWithConfigEntry):
+    """Handle an options flow."""
 
     async def async_step_init(
         self, user_input: dict[str, Any] | None = None
@@ -308,13 +309,13 @@ class FritzBoxToolsOptionsFlowHandler(OptionsFlow):
             {
                 vol.Optional(
                     CONF_CONSIDER_HOME,
-                    default=self.config_entry.options.get(
+                    default=self.options.get(
                         CONF_CONSIDER_HOME, DEFAULT_CONSIDER_HOME.total_seconds()
                     ),
                 ): vol.All(vol.Coerce(int), vol.Clamp(min=0, max=900)),
                 vol.Optional(
                     CONF_OLD_DISCOVERY,
-                    default=self.config_entry.options.get(
+                    default=self.options.get(
                         CONF_OLD_DISCOVERY, DEFAULT_CONF_OLD_DISCOVERY
                     ),
                 ): bool,
diff --git a/tests/components/fritz/test_config_flow.py b/tests/components/fritz/test_config_flow.py
index ded7cda0dea..5b87d897dd9 100644
--- a/tests/components/fritz/test_config_flow.py
+++ b/tests/components/fritz/test_config_flow.py
@@ -9,11 +9,13 @@ from fritzconnection.core.exceptions import (
 )
 import pytest
 
+from homeassistant import data_entry_flow
 from homeassistant.components.device_tracker import (
     CONF_CONSIDER_HOME,
     DEFAULT_CONSIDER_HOME,
 )
 from homeassistant.components.fritz.const import (
+    CONF_OLD_DISCOVERY,
     DOMAIN,
     ERROR_AUTH_INVALID,
     ERROR_CANNOT_CONNECT,
@@ -453,28 +455,24 @@ async def test_ssdp_exception(hass: HomeAssistant, mock_get_source_ip) -> None:
         assert result["step_id"] == "confirm"
 
 
-async def test_options_flow(
-    hass: HomeAssistant, fc_class_mock, mock_get_source_ip
-) -> None:
+async def test_options_flow(hass: HomeAssistant) -> None:
     """Test options flow."""
 
     mock_config = MockConfigEntry(domain=DOMAIN, data=MOCK_USER_DATA)
     mock_config.add_to_hass(hass)
 
-    with patch(
-        "homeassistant.components.fritz.config_flow.FritzConnection",
-        side_effect=fc_class_mock,
-    ), patch("homeassistant.components.fritz.common.FritzBoxTools"):
-        result = await hass.config_entries.options.async_init(mock_config.entry_id)
-        assert result["type"] == FlowResultType.FORM
-        assert result["step_id"] == "init"
+    result = await hass.config_entries.options.async_init(mock_config.entry_id)
+    await hass.async_block_till_done()
+    result = await hass.config_entries.options.async_configure(
+        result["flow_id"],
+        user_input={
+            CONF_CONSIDER_HOME: 37,
+        },
+    )
+    await hass.async_block_till_done()
 
-        result = await hass.config_entries.options.async_init(mock_config.entry_id)
-        result = await hass.config_entries.options.async_configure(
-            result["flow_id"],
-            user_input={
-                CONF_CONSIDER_HOME: 37,
-            },
-        )
-        assert result["type"] == FlowResultType.CREATE_ENTRY
-        assert mock_config.options[CONF_CONSIDER_HOME] == 37
+    assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
+    assert result["data"] == {
+        CONF_OLD_DISCOVERY: False,
+        CONF_CONSIDER_HOME: 37,
+    }