From cc14dfa31cca60276c9c5a689b9521d60c701c94 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Sun, 26 Apr 2020 23:35:04 -0700
Subject: [PATCH] Allow ignoring discovery config flow helper (#34740)

---
 homeassistant/helpers/config_entry_flow.py |  9 ++---
 tests/helpers/test_config_entry_flow.py    | 38 +++++++++++++++++++++-
 2 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/homeassistant/helpers/config_entry_flow.py b/homeassistant/helpers/config_entry_flow.py
index 323c6907411..81881d943cd 100644
--- a/homeassistant/helpers/config_entry_flow.py
+++ b/homeassistant/helpers/config_entry_flow.py
@@ -33,6 +33,8 @@ class DiscoveryFlowHandler(config_entries.ConfigFlow):
         if self._async_current_entries():
             return self.async_abort(reason="single_instance_allowed")
 
+        await self.async_set_unique_id(self._domain, raise_on_progress=False)
+
         return await self.async_step_confirm()
 
     async def async_step_confirm(self, user_input=None):
@@ -40,10 +42,7 @@ class DiscoveryFlowHandler(config_entries.ConfigFlow):
         if user_input is None:
             return self.async_show_form(step_id="confirm")
 
-        if (  # pylint: disable=no-member # https://github.com/PyCQA/pylint/issues/3167
-            self.context
-            and self.context.get("source") != config_entries.SOURCE_DISCOVERY
-        ):
+        if self.source == config_entries.SOURCE_USER:
             # Get current discovered entries.
             in_progress = self._async_in_progress()
 
@@ -67,6 +66,8 @@ class DiscoveryFlowHandler(config_entries.ConfigFlow):
         if self._async_in_progress() or self._async_current_entries():
             return self.async_abort(reason="single_instance_allowed")
 
+        await self.async_set_unique_id(self._domain)
+
         return await self.async_step_confirm()
 
     async_step_zeroconf = async_step_discovery
diff --git a/tests/helpers/test_config_entry_flow.py b/tests/helpers/test_config_entry_flow.py
index 3126c2c0c8d..868ff082ec3 100644
--- a/tests/helpers/test_config_entry_flow.py
+++ b/tests/helpers/test_config_entry_flow.py
@@ -44,6 +44,7 @@ async def test_single_entry_allowed(hass, discovery_flow_conf):
     """Test only a single entry is allowed."""
     flow = config_entries.HANDLERS["test"]()
     flow.hass = hass
+    flow.context = {}
 
     MockConfigEntry(domain="test").add_to_hass(hass)
     result = await flow.async_step_user()
@@ -67,6 +68,7 @@ async def test_user_has_confirmation(hass, discovery_flow_conf):
     """Test user requires no confirmation to setup."""
     flow = config_entries.HANDLERS["test"]()
     flow.hass = hass
+    flow.context = {}
     discovery_flow_conf["discovered"] = True
 
     result = await flow.async_step_user()
@@ -93,7 +95,7 @@ async def test_discovery_confirmation(hass, discovery_flow_conf, source):
     """Test we ask for confirmation via discovery."""
     flow = config_entries.HANDLERS["test"]()
     flow.hass = hass
-    flow.context = {}
+    flow.context = {"source": source}
 
     result = await getattr(flow, f"async_step_{source}")({})
 
@@ -150,6 +152,7 @@ async def test_import_no_confirmation(hass, discovery_flow_conf):
     """Test import requires no confirmation to set up."""
     flow = config_entries.HANDLERS["test"]()
     flow.hass = hass
+    flow.context = {}
     discovery_flow_conf["discovered"] = True
 
     result = await flow.async_step_import(None)
@@ -160,6 +163,7 @@ async def test_import_single_instance(hass, discovery_flow_conf):
     """Test import doesn't create second instance."""
     flow = config_entries.HANDLERS["test"]()
     flow.hass = hass
+    flow.context = {}
     discovery_flow_conf["discovered"] = True
     MockConfigEntry(domain="test").add_to_hass(hass)
 
@@ -167,6 +171,38 @@ async def test_import_single_instance(hass, discovery_flow_conf):
     assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
 
 
+async def test_ignored_discoveries(hass, discovery_flow_conf):
+    """Test we can ignore discovered entries."""
+    mock_entity_platform(hass, "config_flow.test", None)
+
+    result = await hass.config_entries.flow.async_init(
+        "test", context={"source": config_entries.SOURCE_DISCOVERY}, data={}
+    )
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+
+    flow = next(
+        (
+            flw
+            for flw in hass.config_entries.flow.async_progress()
+            if flw["flow_id"] == result["flow_id"]
+        ),
+        None,
+    )
+
+    # Ignore it.
+    await hass.config_entries.flow.async_init(
+        flow["handler"],
+        context={"source": config_entries.SOURCE_IGNORE},
+        data={"unique_id": flow["context"]["unique_id"]},
+    )
+
+    # Second discovery should be aborted
+    result = await hass.config_entries.flow.async_init(
+        "test", context={"source": config_entries.SOURCE_DISCOVERY}, data={}
+    )
+    assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
+
+
 async def test_webhook_single_entry_allowed(hass, webhook_flow_conf):
     """Test only a single entry is allowed."""
     flow = config_entries.HANDLERS["test_single"]()