From 4aed1521b5064f3874a6ebf88527979af7e15e49 Mon Sep 17 00:00:00 2001
From: Jason Rumney <jasonrumney@gmail.com>
Date: Sun, 15 Dec 2024 16:35:37 +0900
Subject: [PATCH] config_flow: collect the results of possible_matches
 generator.

Instead of trying to turn the possible_matches generator into an async
generator, instead collect its results and call the whole thing inside
a hass async task.

Hopefully this gets rid of the final warnings about I/O being done in
the main thread during config flow.

Issue #2584
---
 custom_components/tuya_local/config_flow.py |  2 +-
 custom_components/tuya_local/device.py      | 16 ++++++++++------
 tests/test_config_flow.py                   |  6 +++---
 tests/test_device.py                        |  1 +
 4 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/custom_components/tuya_local/config_flow.py b/custom_components/tuya_local/config_flow.py
index 061677d80d..bd8080a274 100644
--- a/custom_components/tuya_local/config_flow.py
+++ b/custom_components/tuya_local/config_flow.py
@@ -376,7 +376,7 @@ async def async_step_select_type(self, user_input=None):
         best_match = 0
         best_matching_type = None
 
-        async for type in self.device.async_possible_types():
+        for type in await self.device.async_possible_types():
             types.append(type.config_type)
             q = type.match_quality(
                 self.device._get_cached_state(),
diff --git a/custom_components/tuya_local/device.py b/custom_components/tuya_local/device.py
index 1a3087bc82..f9f7a00edb 100644
--- a/custom_components/tuya_local/device.py
+++ b/custom_components/tuya_local/device.py
@@ -33,6 +33,11 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+def _collect_possible_matches(cached_state, product_ids):
+    """Collect possible matches from generator into an array."""
+    return list(possible_matches(cached_state, product_ids))
+
+
 class TuyaLocalDevice(object):
     def __init__(
         self,
@@ -396,19 +401,18 @@ async def async_possible_types(self):
             await self.async_refresh()
             cached_state = self._get_cached_state()
 
-        for matched in await self._hass.async_add_executor_job(
-            possible_matches,
+        return await self._hass.async_add_executor_job(
+            _collect_possible_matches,
             cached_state,
             self._product_ids,
-        ):
-            await asyncio.sleep(0)
-            yield matched
+        )
 
     async def async_inferred_type(self):
         best_match = None
         best_quality = 0
         cached_state = self._get_cached_state()
-        async for config in self.async_possible_types():
+        possible = await self.async_possible_types()
+        for config in possible:
             quality = config.match_quality(cached_state, self._product_ids)
             _LOGGER.info(
                 "%s considering %s with quality %s",
diff --git a/tests/test_config_flow.py b/tests/test_config_flow.py
index 50a0c26c92..27c35fc4f8 100644
--- a/tests/test_config_flow.py
+++ b/tests/test_config_flow.py
@@ -350,9 +350,9 @@ def setup_device_mock(mock, failure=False, type="test"):
     mock_type.legacy_type = type
     mock_type.config_type = type
     mock_type.match_quality.return_value = 100
-    mock_iter = MagicMock()
-    mock_iter.__aiter__.return_value = [mock_type] if not failure else []
-    mock.async_possible_types = MagicMock(return_value=mock_iter)
+    mock.async_possible_types = AsyncMock(
+        return_value=[mock_type] if not failure else []
+    )
 
 
 @pytest.mark.asyncio
diff --git a/tests/test_device.py b/tests/test_device.py
index 4d57463b33..6937ecb351 100644
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -24,6 +24,7 @@ def setUp(self):
         self.hass().data = {"tuya_local": {}}
 
         def job(func, *args):
+            print(f"{args}")
             return func(*args)
 
         self.hass().async_add_executor_job = AsyncMock()