Store toon flow data in flow handler attributes (#127180)

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()

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,