Fix via_device race in habitica (#177718)

Co-authored-by: Markus Tuominen <3738613+Markus98@users.noreply.github.com>

homeassistant/components/habitica/__init__.py

@@ -1,8 +1,10 @@
 """The habitica integration."""
 
+from typing import TYPE_CHECKING
 from uuid import UUID
 
 from habiticalib import Habitica
+from yarl import URL
 
 from homeassistant.components.notify import DOMAIN as NOTIFY_DOMAIN
 from homeassistant.const import CONF_API_KEY, CONF_URL, CONF_VERIFY_SSL, Platform
@@ -16,7 +18,7 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.util.hass_dict import HassKey
 
-from .const import CONF_API_USER, DOMAIN, X_CLIENT
+from .const import CONF_API_USER, DOMAIN, MANUFACTURER, NAME, X_CLIENT
 from .coordinator import (
     HabiticaConfigEntry,
     HabiticaDataUpdateCoordinator,
@@ -71,6 +73,23 @@ async def async_setup_entry(
 
     config_entry.runtime_data = coordinator
 
+    if TYPE_CHECKING:
+        assert config_entry.unique_id
+
+    # Register the user device so children (party device) can resolve it
+    # deterministically as their via_device parent before platforms are set up.
+    device_reg.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        entry_type=dr.DeviceEntryType.SERVICE,
+        manufacturer=MANUFACTURER,
+        model=NAME,
+        name=coordinator.data.user.profile.name,
+        configuration_url=(
+            URL(config_entry.data[CONF_URL]) / "profile" / config_entry.unique_id
+        ),
+        identifiers={(DOMAIN, config_entry.unique_id)},
+    )
+
     party = coordinator.data.user.party.id
     hass.data.setdefault(HABITICA_KEY, {})
 
@@ -81,6 +100,23 @@ async def async_setup_entry(
         hass.data[HABITICA_KEY][party] = party_coordinator
         party_added_by_this_entry = party
 
+    if party is not None:
+        # Register the party device so party member devices can resolve it
+        # deterministically as their via_device parent.
+        device_reg.async_get_or_create(
+            config_entry_id=config_entry.entry_id,
+            entry_type=dr.DeviceEntryType.SERVICE,
+            manufacturer=MANUFACTURER,
+            model=NAME,
+            name=hass.data[HABITICA_KEY][party].data.party.summary,
+            identifiers={(DOMAIN, f"{config_entry.unique_id}_{party!s}")},
+            via_device_id=dr.async_get_device_id_by_identifier(
+                hass,
+                (DOMAIN, config_entry.unique_id),
+                config_entry_id=config_entry.entry_id,
+            ),
+        )
+
     @callback
     def _party_update_listener() -> None:
         """On party change, unload coordinator, remove device and reload."""

homeassistant/components/habitica/entity.py

@@ -8,6 +8,7 @@ from yarl import URL
 
 from homeassistant.config_entries import ConfigSubentry
 from homeassistant.const import CONF_URL
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -57,14 +58,14 @@ class HabiticaBase(CoordinatorEntity[HabiticaDataUpdateCoordinator]):
         )
 
         if subentry:
-            self._attr_device_info.update(
+            self._attr_device_info["via_device_id"] = (
-                DeviceInfo(
+                dr.async_get_device_id_by_identifier(
-                    via_device=(
+                    coordinator.hass,
-                        (
+                    (
-                            DOMAIN,
+                        DOMAIN,
-                            f"{coordinator.config_entry.unique_id}_{self.user.party.id}",
+                        f"{coordinator.config_entry.unique_id}_{self.user.party.id}",
-                        )
+                    ),
-                    )
+                    config_entry_id=coordinator.config_entry.entry_id,
                 )
             )
 
@@ -138,6 +139,10 @@ class HabiticaPartyBase(CoordinatorEntity[HabiticaPartyCoordinator]):
             model=NAME,
             name=coordinator.data.party.summary,
             identifiers={(DOMAIN, unique_id)},
-            via_device=(DOMAIN, config_entry.unique_id),
+            via_device_id=dr.async_get_device_id_by_identifier(
+                coordinator.hass,
+                (DOMAIN, config_entry.unique_id),
+                config_entry_id=config_entry.entry_id,
+            ),
         )
         self.content = content

tests/components/habitica/test_init.py

@@ -179,3 +179,40 @@ async def test_remove_party_and_reload(
         hass.states.get("notify.test_user_private_message_test_partymember_displayname")
         is None
     )
+
+
+@pytest.mark.usefixtures("habitica")
+async def test_device_via_device_links(
+    hass: HomeAssistant,
+    config_entry_with_subentry: MockConfigEntry,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test the via_device links between user, party and party member devices."""
+    group_id = "1e87097c-4c03-4f8c-a475-67cc7da7f409"
+    member_id = "ffce870c-3ff3-4fa4-bad1-87612e52b8e7"
+
+    config_entry_with_subentry.add_to_hass(hass)
+    await hass.config_entries.async_setup(config_entry_with_subentry.entry_id)
+
+    assert config_entry_with_subentry.state is ConfigEntryState.LOADED
+
+    unique_id = config_entry_with_subentry.unique_id
+    entry_id = config_entry_with_subentry.entry_id
+
+    user_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, unique_id), entry_id
+    )
+    assert user_device is not None
+    assert user_device.via_device_id is None
+
+    party_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, f"{unique_id}_{group_id}"), entry_id
+    )
+    assert party_device is not None
+    assert party_device.via_device_id == user_device.id
+
+    member_device = device_registry.async_get_device_by_identifier(
+        (DOMAIN, member_id), entry_id
+    )
+    assert member_device is not None
+    assert member_device.via_device_id == party_device.id