Add optional context parameter to async_start_reauth (#76077)

homeassistant/components/xiaomi_ble/sensor.py

@@ -181,25 +181,7 @@ def process_service_info(
         and data.encryption_scheme != EncryptionScheme.NONE
         and not data.bindkey_verified
     ):
-        flow_context = {
-            "source": config_entries.SOURCE_REAUTH,
-            "entry_id": entry.entry_id,
-            "title_placeholders": {"name": entry.title},
-            "unique_id": entry.unique_id,
-            "device": data,
-        }
-
-        for flow in hass.config_entries.flow.async_progress_by_handler(DOMAIN):
-            if flow["context"] == flow_context:
-                break
-        else:
-            hass.async_create_task(
-                hass.config_entries.flow.async_init(
-                    DOMAIN,
-                    context=flow_context,
-                    data=entry.data,
-                )
-            )
+        entry.async_start_reauth(hass, context={"device": data})
 
     return sensor_update_to_bluetooth_data_update(update)
 

homeassistant/config_entries.py

@@ -640,7 +640,9 @@ class ConfigEntry:
                 await asyncio.gather(*pending)
 
     @callback
-    def async_start_reauth(self, hass: HomeAssistant) -> None:
+    def async_start_reauth(
+        self, hass: HomeAssistant, context: dict[str, Any] | None = None
+    ) -> None:
         """Start a reauth flow."""
         flow_context = {
             "source": SOURCE_REAUTH,
@@ -649,6 +651,9 @@ class ConfigEntry:
             "unique_id": self.unique_id,
         }
 
+        if context:
+            flow_context.update(context)
+
         for flow in hass.config_entries.flow.async_progress_by_handler(self.domain):
             if flow["context"] == flow_context:
                 return

tests/test_config_entries.py

@@ -3273,3 +3273,30 @@ async def test_disallow_entry_reload_with_setup_in_progresss(hass, manager):
     with pytest.raises(config_entries.OperationNotAllowed):
         assert await manager.async_reload(entry.entry_id)
     assert entry.state is config_entries.ConfigEntryState.SETUP_IN_PROGRESS
+
+
+async def test_reauth(hass):
+    """Test the async_reauth_helper."""
+    entry = MockConfigEntry(title="test_title", domain="test")
+
+    mock_setup_entry = AsyncMock(return_value=True)
+    mock_integration(hass, MockModule("test", async_setup_entry=mock_setup_entry))
+    mock_entity_platform(hass, "config_flow.test", None)
+
+    await entry.async_setup(hass)
+    await hass.async_block_till_done()
+
+    entry.async_start_reauth(hass, {"extra_context": "some_extra_context"})
+    await hass.async_block_till_done()
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+    assert flows[0]["context"]["entry_id"] == entry.entry_id
+    assert flows[0]["context"]["source"] == config_entries.SOURCE_REAUTH
+    assert flows[0]["context"]["title_placeholders"] == {"name": "test_title"}
+    assert flows[0]["context"]["extra_context"] == "some_extra_context"
+
+    # Check we can't start duplicate flows
+    entry.async_start_reauth(hass, {"extra_context": "some_extra_context"})
+    await hass.async_block_till_done()
+    assert len(flows) == 1