Fix setup timings when config entry platform loads are not awaited (#113959)

* Move setup time logging into the context manager

We were fetching the time twice but since the context
manager already has the timing, move it there

* remove log setup assertions from integration test

* tweak logging to give us better data for tracking issues

* redundant

* adjust

* preen

* fixes

* adjust

* make api change internal so nobody uses it

* coverage

* fix test

* fix more tests

* coverage

* more tests assuming internal calls

* fix more

* adjust

* adjust

* fix axis tests

* fix broadlink -- it does not call async_forward_entry_setup

* missed some

* remove useless patch

* rename, detect it both ways

* clear

* debug

* try to fix

* handle phase finishing out while paused

* where its set does not need to know its late as that is an implemenation detail of setup

* where its set does not need to know its late as that is an implemenation detail of setup

* tweak

* simplify

* reduce complexity

* revert order change as it makes review harder

* revert naming changes as it makes review harder

* improve comment

* improve debug

* late dispatch test

* test the other way as well

* Update setup.py

* Update setup.py

* Update setup.py

* simplify

* reduce

tests/test_setup.py

@@ -4,6 +4,7 @@ import asyncio
 import threading
 from unittest.mock import ANY, AsyncMock, Mock, patch
 
+from freezegun.api import FrozenDateTimeFactory
 import pytest
 import voluptuous as vol
 
@@ -16,6 +17,10 @@ from homeassistant.helpers.config_validation import (
     PLATFORM_SCHEMA,
     PLATFORM_SCHEMA_BASE,
 )
+from homeassistant.helpers.dispatcher import (
+    async_dispatcher_connect,
+    async_dispatcher_send,
+)
 
 from .common import (
     MockConfigEntry,
@@ -739,7 +744,9 @@ async def test_async_start_setup_running(hass: HomeAssistant) -> None:
         assert not setup_started
 
 
-async def test_async_start_setup_config_entry(hass: HomeAssistant) -> None:
+async def test_async_start_setup_config_entry(
+    hass: HomeAssistant, freezer: FrozenDateTimeFactory
+) -> None:
     """Test setup started keeps track of setup times with a config entry."""
     hass.set_state(CoreState.not_running)
     setup_started: dict[tuple[str, str | None], float]
@@ -778,6 +785,7 @@ async def test_async_start_setup_config_entry(hass: HomeAssistant) -> None:
         phase=setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP,
     ):
         assert isinstance(setup_started[("august", "entry_id")], float)
+
     # Platforms outside of CONFIG_ENTRY_SETUP should be tracked
     # This simulates a late platform forward
     assert setup_time["august"] == {
@@ -788,6 +796,38 @@ async def test_async_start_setup_config_entry(hass: HomeAssistant) -> None:
         },
     }
 
+    shorter_time = setup_time["august"]["entry_id"][
+        setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP
+    ]
+    # Setup another platform, but make it take longer
+    with setup.async_start_setup(
+        hass,
+        integration="august",
+        group="entry_id",
+        phase=setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP,
+    ):
+        freezer.tick(10)
+        assert isinstance(setup_started[("august", "entry_id")], float)
+
+    longer_time = setup_time["august"]["entry_id"][
+        setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP
+    ]
+    assert longer_time > shorter_time
+    # Setup another platform, but make it take shorter
+    with setup.async_start_setup(
+        hass,
+        integration="august",
+        group="entry_id",
+        phase=setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP,
+    ):
+        assert isinstance(setup_started[("august", "entry_id")], float)
+
+    # Ensure we keep the longest time
+    assert (
+        setup_time["august"]["entry_id"][setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP]
+        == longer_time
+    )
+
     with setup.async_start_setup(
         hass,
         integration="august",
@@ -815,6 +855,106 @@ async def test_async_start_setup_config_entry(hass: HomeAssistant) -> None:
     }
 
 
+async def test_async_start_setup_config_entry_late_platform(
+    hass: HomeAssistant, freezer: FrozenDateTimeFactory
+) -> None:
+    """Test setup started tracks config entry time with a late platform load."""
+    hass.set_state(CoreState.not_running)
+    setup_started: dict[tuple[str, str | None], float]
+    setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {})
+    setup_time = setup._setup_times(hass)
+
+    with setup.async_start_setup(
+        hass, integration="august", phase=setup.SetupPhases.SETUP
+    ):
+        freezer.tick(10)
+        assert isinstance(setup_started[("august", None)], float)
+
+    with setup.async_start_setup(
+        hass,
+        integration="august",
+        group="entry_id",
+        phase=setup.SetupPhases.CONFIG_ENTRY_SETUP,
+    ):
+        assert isinstance(setup_started[("august", "entry_id")], float)
+
+        @callback
+        def async_late_platform_load():
+            with setup.async_pause_setup(hass, setup.SetupPhases.WAIT_IMPORT_PLATFORMS):
+                freezer.tick(100)
+            with setup.async_start_setup(
+                hass,
+                integration="august",
+                group="entry_id",
+                phase=setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP,
+            ):
+                freezer.tick(20)
+                assert isinstance(setup_started[("august", "entry_id")], float)
+
+        disconnect = async_dispatcher_connect(
+            hass, "late_platform_load_test", async_late_platform_load
+        )
+
+    # Dispatch a late platform load
+    async_dispatcher_send(hass, "late_platform_load_test")
+    disconnect()
+
+    # CONFIG_ENTRY_PLATFORM_SETUP is late dispatched, so it should be tracked
+    # but any waiting time should not be because it's blocking the setup
+    assert setup_time["august"] == {
+        None: {setup.SetupPhases.SETUP: 10.0},
+        "entry_id": {
+            setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP: 20.0,
+            setup.SetupPhases.CONFIG_ENTRY_SETUP: 0.0,
+        },
+    }
+
+
+async def test_async_start_setup_config_entry_platform_wait(
+    hass: HomeAssistant, freezer: FrozenDateTimeFactory
+) -> None:
+    """Test setup started tracks wait time when a platform loads inside of config entry setup."""
+    hass.set_state(CoreState.not_running)
+    setup_started: dict[tuple[str, str | None], float]
+    setup_started = hass.data.setdefault(setup.DATA_SETUP_STARTED, {})
+    setup_time = setup._setup_times(hass)
+
+    with setup.async_start_setup(
+        hass, integration="august", phase=setup.SetupPhases.SETUP
+    ):
+        freezer.tick(10)
+        assert isinstance(setup_started[("august", None)], float)
+
+    with setup.async_start_setup(
+        hass,
+        integration="august",
+        group="entry_id",
+        phase=setup.SetupPhases.CONFIG_ENTRY_SETUP,
+    ):
+        assert isinstance(setup_started[("august", "entry_id")], float)
+
+        with setup.async_pause_setup(hass, setup.SetupPhases.WAIT_IMPORT_PLATFORMS):
+            freezer.tick(100)
+        with setup.async_start_setup(
+            hass,
+            integration="august",
+            group="entry_id",
+            phase=setup.SetupPhases.CONFIG_ENTRY_PLATFORM_SETUP,
+        ):
+            freezer.tick(20)
+            assert isinstance(setup_started[("august", "entry_id")], float)
+
+    # CONFIG_ENTRY_PLATFORM_SETUP is run inside of CONFIG_ENTRY_SETUP, so it should not
+    # be tracked, but any wait time should still be tracked because its blocking the setup
+    assert setup_time["august"] == {
+        None: {setup.SetupPhases.SETUP: 10.0},
+        "entry_id": {
+            setup.SetupPhases.WAIT_IMPORT_PLATFORMS: -100.0,
+            setup.SetupPhases.CONFIG_ENTRY_SETUP: 120.0,
+        },
+    }
+
+
 async def test_async_start_setup_top_level_yaml(hass: HomeAssistant) -> None:
     """Test setup started context manager keeps track of setup times with modern yaml."""
     hass.set_state(CoreState.not_running)