Repository: home-assistant/core
File: tests/test_config_entries.py (Python)

Code before:
    hass: HomeAssistant,
    manager: config_entries.ConfigEntries,
) -> None:
    """Test that discovery flows are marked dismiss protected when user configures."""
    mock_integration(
        hass,
        MockModule("comp", async_setup_entry=AsyncMock(return_value=True)),
    )
    mock_platform(hass, "comp.config_flow", None)

    class TestFlow(config_entries.ConfigFlow):
        """Test flow."""

        VERSION = 1

        async def async_step_zeroconf(self, discovery_info):
            """Test zeroconf step."""
            return self.async_show_form(step_id="confirm")

        async def async_step_confirm(self, user_input=None):
            """Test confirm step."""
            if user_input is None:
                return self.async_show_form(step_id="confirm")
            return self.async_create_entry(title="test", data={})

    def _get_flow_context(flow_id: str) -> config_entries.ConfigFlowContext:
        """Get the flow context from async_progress."""
        return next(
            flow["context"]
            for flow in manager.flow.async_progress()
            if flow["flow_id"] == flow_id
        )

    with mock_config_flow("comp", TestFlow):
        result = await manager.flow.async_init(
            "comp",
            context={"source": config_entries.SOURCE_ZEROCONF},
            data=ZeroconfServiceInfo(
                ip_address=ip_address("192.168.1.1"),
                ip_addresses=[ip_address("192.168.1.1")],
                hostname="test.local.",
                name="test._tcp.local.",
                port=80,
                properties={},
                type="_tcp.local.",
            ),
        )
        assert result["type"] == data_entry_flow.FlowResultType.FORM
        assert result["step_id"] == "confirm"

        # Before user interaction, dismiss_protected should not be set

Review comment (refactor):
I think we can make `_get_flow_context` a test helper in this module by adding `manager` as a parameter in the function signature.

Code after:
    task = hass.async_create_task(manager.async_setup(entry.entry_id))
    await start.wait()
    task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await task

    abort.set()


async def test_discovery_flow_dismiss_protected_on_configure(
    hass: HomeAssistant,
    manager: config_entries.ConfigEntries,
) -> None:
    """Test that discovery flows are marked dismiss protected when user configures."""
    mock_integration(
        hass,
        MockModule("comp", async_setup_entry=AsyncMock(return_value=True)),
    )
    mock_platform(hass, "comp.config_flow", None)

    class TestFlow(config_entries.ConfigFlow):
        """Test flow."""

        VERSION = 1

        async def async_step_zeroconf(self, discovery_info):
            """Test zeroconf step."""
            return self.async_show_form(step_id="confirm")

        async def async_step_confirm(self, user_input=None):
            """Test confirm step."""
            if user_input is None:
                return self.async_show_form(step_id="confirm")
            return self.async_create_entry(title="test", data={})

    with mock_config_flow("comp", TestFlow):
        result = await manager.flow.async_init(
            "comp",
            context={"source": config_entries.SOURCE_ZEROCONF},
            data=ZeroconfServiceInfo(
                ip_address=ip_address("192.168.1.1"),
                ip_addresses=[ip_address("192.168.1.1")],
                hostname="test.local.",
                name="test._tcp.local.",
                port=80,
                properties={},
                type="_tcp.local.",
            ),
        )
        assert result["type"] == data_entry_flow.FlowResultType.FORM