From d9bba25f67a7c82b8dc536668c1895743335ab6f Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Tue, 1 Oct 2024 11:15:54 +0200 Subject: [PATCH] Store toon flow data in flow handler attributes (#127180) --- homeassistant/components/toon/config_flow.py | 7 ++++--- tests/components/toon/test_config_flow.py | 5 +++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/homeassistant/components/toon/config_flow.py b/homeassistant/components/toon/config_flow.py index af9f7b06850..450d2472a6c 100644 --- a/homeassistant/components/toon/config_flow.py +++ b/homeassistant/components/toon/config_flow.py @@ -23,6 +23,7 @@ class ToonFlowHandler(AbstractOAuth2FlowHandler, domain=DOMAIN): agreements: list[Agreement] data: dict[str, Any] + migrate_entry: str | None = None @property def logger(self) -> logging.Logger: @@ -58,7 +59,7 @@ class ToonFlowHandler(AbstractOAuth2FlowHandler, domain=DOMAIN): """ if import_data is not None and CONF_MIGRATE in import_data: - self.context.update({CONF_MIGRATE: import_data[CONF_MIGRATE]}) + self.migrate_entry = import_data[CONF_MIGRATE] else: await self._async_handle_discovery_without_unique_id() @@ -88,8 +89,8 @@ class ToonFlowHandler(AbstractOAuth2FlowHandler, domain=DOMAIN): return await self._create_entry(self.agreements[agreement_index]) async def _create_entry(self, agreement: Agreement) -> ConfigFlowResult: - if CONF_MIGRATE in self.context: - await self.hass.config_entries.async_remove(self.context[CONF_MIGRATE]) + if self.migrate_entry: + await self.hass.config_entries.async_remove(self.migrate_entry) await self.async_set_unique_id(agreement.agreement_id) self._abort_if_unique_id_configured() diff --git a/tests/components/toon/test_config_flow.py b/tests/components/toon/test_config_flow.py index 492e2a220ad..228cb0b0239 100644 --- a/tests/components/toon/test_config_flow.py +++ b/tests/components/toon/test_config_flow.py @@ -6,7 +6,7 @@ from unittest.mock import patch import pytest from toonapi import Agreement, ToonError -from homeassistant.components.toon.const import CONF_AGREEMENT, CONF_MIGRATE, DOMAIN +from homeassistant.components.toon.const import CONF_AGREEMENT, DOMAIN from homeassistant.config import async_process_ha_core_config from homeassistant.config_entries import SOURCE_IMPORT, SOURCE_USER from homeassistant.const import CONF_CLIENT_ID, CONF_CLIENT_SECRET @@ -324,7 +324,8 @@ async def test_import_migration( flows = hass.config_entries.flow.async_progress() assert len(flows) == 1 - assert flows[0]["context"][CONF_MIGRATE] == old_entry.entry_id + flow = hass.config_entries.flow._progress[flows[0]["flow_id"]] + assert flow.migrate_entry == old_entry.entry_id state = config_entry_oauth2_flow._encode_jwt( hass,