Fix race condition on eheimdigital coordinator setup

homeassistant/components/eheimdigital/coordinator.py

@@ -2,16 +2,18 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Callable
 
 from aiohttp import ClientError
 from eheimdigital.device import EheimDigitalDevice
 from eheimdigital.hub import EheimDigitalHub
-from eheimdigital.types import EheimDeviceType
+from eheimdigital.types import EheimDeviceType, EheimDigitalClientError
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.entity_component import DEFAULT_SCAN_INTERVAL
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
@@ -43,12 +45,14 @@
             name=DOMAIN,
             update_interval=DEFAULT_SCAN_INTERVAL,
         )
+        self.main_device_added_event = asyncio.Event()
         self.hub = EheimDigitalHub(
             host=self.config_entry.data[CONF_HOST],
             session=async_get_clientsession(hass),
             loop=hass.loop,
             receive_callback=self._async_receive_callback,
             device_found_callback=self._async_device_found,
+            main_device_added_event=self.main_device_added_event,
         )
         self.known_devices: set[str] = set()
         self.platform_callbacks: set[AsyncSetupDeviceEntitiesCallback] = set()
@@ -76,8 +80,17 @@
         self.async_set_updated_data(self.hub.devices)
 
     async def _async_setup(self) -> None:
-        await self.hub.connect()
-        await self.hub.update()
+        try:
+            await self.hub.connect()
+            async with asyncio.timeout(2):
+                # This event gets triggered when the first message is received from
+                # the device, it contains the data necessary to create the main device.
+                # This removes the race condition where the main device is accessed
+                # before the response from the device is parsed.
+                await self.main_device_added_event.wait()
+            await self.hub.update()
+        except (TimeoutError, EheimDigitalClientError) as err:
+            raise ConfigEntryNotReady from err
 
     async def _async_update_data(self) -> dict[str, EheimDigitalDevice]:
         try:

tests/components/eheimdigital/conftest.py

@@ -11,6 +11,7 @@
 
 from homeassistant.components.eheimdigital.const import DOMAIN
 from homeassistant.const import CONF_HOST
+from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry
 
@@ -79,3 +80,15 @@ def eheimdigital_hub_mock(
         }
         eheimdigital_hub_mock.return_value.main = classic_led_ctrl_mock
         yield eheimdigital_hub_mock
+
+
+async def init_integration(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
+    """Initialize the integration."""
+
+    mock_config_entry.add_to_hass(hass)
+    with patch(
+        "homeassistant.components.eheimdigital.coordinator.asyncio.Event", new=AsyncMock
+    ):
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)

tests/components/eheimdigital/test_climate.py

@@ -1,6 +1,6 @@
 """Tests for the climate module."""
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from eheimdigital.types import (
     EheimDeviceType,
@@ -31,6 +31,8 @@
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
 
+from .conftest import init_integration
+
 from tests.common import MockConfigEntry, snapshot_platform
 
 
@@ -45,7 +47,13 @@ async def test_setup_heater(
     """Test climate platform setup for heater."""
     mock_config_entry.add_to_hass(hass)
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.CLIMATE]):
+    with (
+        patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.CLIMATE]),
+        patch(
+            "homeassistant.components.eheimdigital.coordinator.asyncio.Event",
+            new=AsyncMock,
+        ),
+    ):
         await hass.config_entries.async_setup(mock_config_entry.entry_id)
 
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
@@ -69,7 +77,13 @@ async def test_dynamic_new_devices(
 
     eheimdigital_hub_mock.return_value.devices = {}
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.CLIMATE]):
+    with (
+        patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.CLIMATE]),
+        patch(
+            "homeassistant.components.eheimdigital.coordinator.asyncio.Event",
+            new=AsyncMock,
+        ),
+    ):
         await hass.config_entries.async_setup(mock_config_entry.entry_id)
 
     assert (
@@ -108,9 +122,7 @@ async def test_set_preset_mode(
     heater_mode: HeaterMode,
 ) -> None:
     """Test setting a preset mode."""
-    mock_config_entry.add_to_hass(hass)
-
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await init_integration(hass, mock_config_entry)
 
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:02", EheimDeviceType.VERSION_EHEIM_EXT_HEATER
@@ -146,9 +158,7 @@ async def test_set_temperature(
     mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test setting a preset mode."""
-    mock_config_entry.add_to_hass(hass)
-
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await init_integration(hass, mock_config_entry)
 
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:02", EheimDeviceType.VERSION_EHEIM_EXT_HEATER
@@ -189,9 +199,7 @@ async def test_set_hvac_mode(
     active: bool,
 ) -> None:
     """Test setting a preset mode."""
-    mock_config_entry.add_to_hass(hass)
-
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await init_integration(hass, mock_config_entry)
 
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:02", EheimDeviceType.VERSION_EHEIM_EXT_HEATER
@@ -231,9 +239,8 @@ async def test_state_update(
     heater_mock.is_heating = False
     heater_mock.operation_mode = HeaterMode.BIO
 
-    mock_config_entry.add_to_hass(hass)
+    await init_integration(hass, mock_config_entry)
 
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:02", EheimDeviceType.VERSION_EHEIM_EXT_HEATER
     )

tests/components/eheimdigital/test_init.py

@@ -8,6 +8,8 @@
 from homeassistant.helpers import device_registry as dr
 from homeassistant.setup import async_setup_component
 
+from .conftest import init_integration
+
 from tests.common import MockConfigEntry
 from tests.typing import WebSocketGenerator
 
@@ -21,9 +23,8 @@ async def test_remove_device(
 ) -> None:
     """Test removing a device."""
     assert await async_setup_component(hass, "config", {})
-    mock_config_entry.add_to_hass(hass)
 
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await init_integration(hass, mock_config_entry)
 
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:01", EheimDeviceType.VERSION_EHEIM_CLASSIC_LED_CTRL_PLUS_E

tests/components/eheimdigital/test_light.py

@@ -1,7 +1,7 @@
 """Tests for the light module."""
 
 from datetime import timedelta
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from aiohttp import ClientError
 from eheimdigital.types import EheimDeviceType, LightMode
@@ -26,6 +26,8 @@
 from homeassistant.helpers import entity_registry as er
 from homeassistant.util.color import value_to_brightness
 
+from .conftest import init_integration
+
 from tests.common import MockConfigEntry, async_fire_time_changed, snapshot_platform
 
 
@@ -51,7 +53,13 @@ async def test_setup_classic_led_ctrl(
 
     classic_led_ctrl_mock.tankconfig = tankconfig
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]):
+    with (
+        patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]),
+        patch(
+            "homeassistant.components.eheimdigital.coordinator.asyncio.Event",
+            new=AsyncMock,
+        ),
+    ):
         await hass.config_entries.async_setup(mock_config_entry.entry_id)
 
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
@@ -75,7 +83,13 @@ async def test_dynamic_new_devices(
 
     eheimdigital_hub_mock.return_value.devices = {}
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]):
+    with (
+        patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]),
+        patch(
+            "homeassistant.components.eheimdigital.coordinator.asyncio.Event",
+            new=AsyncMock,
+        ),
+    ):
         await hass.config_entries.async_setup(mock_config_entry.entry_id)
 
     assert (
@@ -106,10 +120,8 @@ async def test_turn_off(
     classic_led_ctrl_mock: MagicMock,
 ) -> None:
     """Test turning off the light."""
-    mock_config_entry.add_to_hass(hass)
+    await init_integration(hass, mock_config_entry)
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]):
-        await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await mock_config_entry.runtime_data._async_device_found(
         "00:00:00:00:00:01", EheimDeviceType.VERSION_EHEIM_CLASSIC_LED_CTRL_PLUS_E
     )
@@ -143,10 +155,8 @@ async def test_turn_on_brightness(
     expected_dim_value: int,
 ) -> None:
     """Test turning on the light with different brightness values."""
-    mock_config_entry.add_to_hass(hass)
+    await init_integration(hass, mock_config_entry)
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]):
-        await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:01", EheimDeviceType.VERSION_EHEIM_CLASSIC_LED_CTRL_PLUS_E
     )
@@ -173,12 +183,10 @@ async def test_turn_on_effect(
     classic_led_ctrl_mock: MagicMock,
 ) -> None:
     """Test turning on the light with an effect value."""
-    mock_config_entry.add_to_hass(hass)
-
     classic_led_ctrl_mock.light_mode = LightMode.MAN_MODE
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]):
-        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await init_integration(hass, mock_config_entry)
+
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:01", EheimDeviceType.VERSION_EHEIM_CLASSIC_LED_CTRL_PLUS_E
     )
@@ -204,10 +212,8 @@ async def test_state_update(
     classic_led_ctrl_mock: MagicMock,
 ) -> None:
     """Test the light state update."""
-    mock_config_entry.add_to_hass(hass)
+    await init_integration(hass, mock_config_entry)
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]):
-        await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:01", EheimDeviceType.VERSION_EHEIM_CLASSIC_LED_CTRL_PLUS_E
     )
@@ -228,10 +234,8 @@ async def test_update_failed(
     freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test an failed update."""
-    mock_config_entry.add_to_hass(hass)
+    await init_integration(hass, mock_config_entry)
 
-    with patch("homeassistant.components.eheimdigital.PLATFORMS", [Platform.LIGHT]):
-        await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await eheimdigital_hub_mock.call_args.kwargs["device_found_callback"](
         "00:00:00:00:00:01", EheimDeviceType.VERSION_EHEIM_CLASSIC_LED_CTRL_PLUS_E
     )