Improve energyid config flow tests (#166749)

tests/components/energyid/test_config_flow.py

@@ -40,6 +40,15 @@ def mock_polling_interval_fixture() -> Generator[int]:
         yield polling_interval
 
 
+@pytest.fixture(autouse=True)
+def mock_setup_entry() -> Generator[AsyncMock]:
+    """Mock setting up a config entry."""
+    with patch(
+        "homeassistant.components.energyid.async_setup_entry", return_value=True
+    ) as mock_setup:
+        yield mock_setup
+
+
 async def test_config_flow_user_step_success_claimed(hass: HomeAssistant) -> None:
     """Test user step where device is already claimed."""
     mock_client = MagicMock()
@@ -105,7 +114,6 @@ async def test_config_flow_auth_and_claim_step_success(hass: HomeAssistant) -> N
             "homeassistant.components.energyid.config_flow.WebhookClient",
             side_effect=mock_webhook_client,
         ),
-        patch("homeassistant.components.energyid.config_flow.asyncio.sleep"),
     ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -146,13 +154,13 @@ async def test_config_flow_claim_timeout(hass: HomeAssistant) -> None:
             "homeassistant.components.energyid.config_flow.WebhookClient",
             return_value=mock_unclaimed_client,
         ),
-        patch(
-            "homeassistant.components.energyid.config_flow.asyncio.sleep",
-        ) as mock_sleep,
     ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_USER}
         )
+
+        assert mock_unclaimed_client.authenticate.call_count == 0
+
         result_external = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             {
@@ -160,8 +168,17 @@ async def test_config_flow_claim_timeout(hass: HomeAssistant) -> None:
                 CONF_PROVISIONING_SECRET: TEST_PROVISIONING_SECRET,
             },
         )
+
         assert result_external["type"] is FlowResultType.EXTERNAL_STEP
 
+        # Wait for the polling to time out.
+        await hass.async_block_till_done()
+
+        # Verify polling actually ran the expected number of times
+        # +1 for the initial attempt before polling starts
+        assert mock_unclaimed_client.authenticate.call_count == MAX_POLLING_ATTEMPTS + 1
+        mock_unclaimed_client.authenticate.reset_mock()
+
         # Simulate polling timeout, then user continuing the flow
         result_after_timeout = await hass.config_entries.flow.async_configure(
             result_external["flow_id"]
@@ -169,14 +186,10 @@ async def test_config_flow_claim_timeout(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
         # After timeout, polling stops and user continues - should see external step again
+        assert mock_unclaimed_client.authenticate.call_count == 1
         assert result_after_timeout["type"] is FlowResultType.EXTERNAL_STEP
         assert result_after_timeout["step_id"] == "auth_and_claim"
 
-        # Verify polling actually ran the expected number of times
-        # Sleep happens at beginning of polling loop, so MAX_POLLING_ATTEMPTS + 1 sleeps
-        # but only MAX_POLLING_ATTEMPTS authentication attempts
-        assert mock_sleep.call_count == MAX_POLLING_ATTEMPTS + 1
-
 
 async def test_duplicate_unique_id_prevented(hass: HomeAssistant) -> None:
     """Test that duplicate device_id (unique_id) is detected and aborted."""
@@ -546,7 +559,6 @@ async def test_config_flow_reauth_needs_claim(hass: HomeAssistant) -> None:
             "homeassistant.components.energyid.config_flow.WebhookClient",
             return_value=mock_client,
         ),
-        patch("homeassistant.components.energyid.config_flow.asyncio.sleep"),
     ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN,
@@ -607,7 +619,6 @@ async def test_polling_stops_on_invalid_auth_error(hass: HomeAssistant) -> None:
             "homeassistant.components.energyid.config_flow.WebhookClient",
             side_effect=mock_webhook_client,
         ),
-        patch("homeassistant.components.energyid.config_flow.asyncio.sleep"),
     ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -655,7 +666,6 @@ async def test_polling_stops_on_cannot_connect_error(hass: HomeAssistant) -> Non
             "homeassistant.components.energyid.config_flow.WebhookClient",
             side_effect=mock_webhook_client,
         ),
-        patch("homeassistant.components.energyid.config_flow.asyncio.sleep"),
     ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -707,7 +717,6 @@ async def test_auth_and_claim_subsequent_auth_error(hass: HomeAssistant) -> None
             "homeassistant.components.energyid.config_flow.WebhookClient",
             side_effect=mock_webhook_client,
         ),
-        patch("homeassistant.components.energyid.config_flow.asyncio.sleep"),
     ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -815,7 +824,6 @@ async def test_polling_cancellation_on_auth_failure(hass: HomeAssistant) -> None
             "homeassistant.components.energyid.config_flow.WebhookClient",
             side_effect=mock_webhook_client,
         ),
-        patch("homeassistant.components.energyid.config_flow.asyncio.sleep"),
     ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -888,7 +896,6 @@ async def test_polling_cancellation_on_success(hass: HomeAssistant) -> None:
             "homeassistant.components.energyid.config_flow.WebhookClient",
             side_effect=mock_webhook_client,
         ),
-        patch("homeassistant.components.energyid.config_flow.asyncio.sleep"),
     ):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": config_entries.SOURCE_USER}
@@ -911,25 +918,16 @@ async def test_polling_cancellation_on_success(hass: HomeAssistant) -> None:
         await hass.async_block_till_done()
 
         # Verify polling made authentication attempt
-        # auth_call_count should be 1 (polling detected device is claimed)
-        assert auth_call_count >= 1
-        claimed_auth_count = auth_call_count
+        assert auth_call_count == 2  # One for polling, one for the final check
 
         # User continues - device is already claimed, polling should be cancelled
         result_done = await hass.config_entries.flow.async_configure(
             result_external["flow_id"]
         )
-        assert result_done["type"] is FlowResultType.EXTERNAL_STEP_DONE
+        assert result_done["type"] is FlowResultType.CREATE_ENTRY
 
-        # Verify polling was cancelled - the auth count should only increase by 1
-        # (for the manual check when user continues, not from polling)
-        assert auth_call_count == claimed_auth_count + 1
-
-        # Final call to create entry
-        final_result = await hass.config_entries.flow.async_configure(
-            result_external["flow_id"]
-        )
-        assert final_result["type"] is FlowResultType.CREATE_ENTRY
+        # Verify polling was cancelled - the auth count should not increase
+        assert auth_call_count == 2
 
         # Wait a bit and verify no further authentication attempts from polling
         await hass.async_block_till_done()