From 55c2bb59c8b852d35b3ae247e1a87b3887c6eb40 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Tue, 30 May 2023 17:45:33 +0200
Subject: [PATCH] Fix race in rfxtrx config flow (#93804)

* Fix race in rfxtrx config flow

* Add timeout

* Use async_timeout.timeout
---
 .../components/rfxtrx/config_flow.py          | 45 ++++++++++++++++---
 1 file changed, 38 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/rfxtrx/config_flow.py b/homeassistant/components/rfxtrx/config_flow.py
index 2aa3bd20b8c..2e054ce4724 100644
--- a/homeassistant/components/rfxtrx/config_flow.py
+++ b/homeassistant/components/rfxtrx/config_flow.py
@@ -2,12 +2,14 @@
 from __future__ import annotations
 
 import asyncio
+from contextlib import suppress
 import copy
 import itertools
 import os
 from typing import Any, TypedDict, cast
 
 import RFXtrx as rfxtrxmod
+from async_timeout import timeout
 import serial
 import serial.tools.list_ports
 import voluptuous as vol
@@ -346,34 +348,57 @@ class OptionsFlow(config_entries.OptionsFlow):
                 entity_migration_map[new_entity_id] = entry
 
         @callback
-        def _handle_state_change(
+        def _handle_state_removed(
             entity_id: str, old_state: State | None, new_state: State | None
         ) -> None:
             # Wait for entities to finish cleanup
-            if new_state is None and entity_id in pending_entities:
-                pending_entities.remove(entity_id)
-            if not pending_entities:
+            if new_state is None and entity_id in entities_to_be_removed:
+                entities_to_be_removed.remove(entity_id)
+            if not entities_to_be_removed:
                 wait_for_entities.set()
 
         # Create a set with entities to be removed which are currently in the state
         # machine
-        pending_entities = {
+        entities_to_be_removed = {
             entry.entity_id
             for entry in entity_migration_map.values()
             if not self.hass.states.async_available(entry.entity_id)
         }
         wait_for_entities = asyncio.Event()
         remove_track_state_changes = async_track_state_change(
-            self.hass, pending_entities, _handle_state_change
+            self.hass, entities_to_be_removed, _handle_state_removed
         )
 
         for entry in entity_migration_map.values():
             entity_registry.async_remove(entry.entity_id)
 
         # Wait for entities to finish cleanup
-        await wait_for_entities.wait()
+        with suppress(asyncio.TimeoutError):
+            async with timeout(10):
+                await wait_for_entities.wait()
         remove_track_state_changes()
 
+        @callback
+        def _handle_state_added(
+            entity_id: str, old_state: State | None, new_state: State | None
+        ) -> None:
+            # Wait for entities to be added
+            if old_state is None and entity_id in entities_to_be_added:
+                entities_to_be_added.remove(entity_id)
+            if not entities_to_be_added:
+                wait_for_entities.set()
+
+        # Create a set with entities to be added to the state machine
+        entities_to_be_added = {
+            entry.entity_id
+            for entry in entity_migration_map.values()
+            if self.hass.states.async_available(entry.entity_id)
+        }
+        wait_for_entities = asyncio.Event()
+        remove_track_state_changes = async_track_state_change(
+            self.hass, entities_to_be_added, _handle_state_added
+        )
+
         for entity_id, entry in entity_migration_map.items():
             entity_registry.async_update_entity(
                 entity_id,
@@ -382,6 +407,12 @@ class OptionsFlow(config_entries.OptionsFlow):
                 icon=entry.icon,
             )
 
+        # Wait for entities to finish renaming
+        with suppress(asyncio.TimeoutError):
+            async with timeout(10):
+                await wait_for_entities.wait()
+        remove_track_state_changes()
+
         device_registry.async_remove_device(old_device)
 
     def _can_add_device(self, new_rfx_obj: rfxtrxmod.RFXtrxEvent) -> bool: