Add via_device permutations to kitchen_sink (#179049)

Co-authored-by: Artur Pragacz <49985303+arturpragacz@users.noreply.github.com>

homeassistant/components/kitchen_sink/device.py

@@ -13,6 +13,7 @@ def async_create_device(
     device_translation_key: str | None,
     device_translation_placeholders: dict[str, str] | None,
     unique_id: str,
+    via_device_id: str | None = None,
 ) -> dr.DeviceEntry:
     """Create a device."""
     device_registry = dr.async_get(hass)
@@ -22,4 +23,5 @@ def async_create_device(
         name=device_name,
         translation_key=device_translation_key,
         translation_placeholders=device_translation_placeholders,
+        via_device_id=via_device_id,
     )

homeassistant/components/kitchen_sink/switch.py

@@ -32,6 +32,24 @@ async def async_setup_entry(
         hass, (DOMAIN, "2_ch_power_strip"), config_entry_id=config_entry.entry_id
     )
 
+    # A hub with a switch device linked to it via its via device.
+    hub = async_create_device(hass, config_entry.entry_id, "Hub", None, None, "hub")
+
+    # A second hub with a power strip linked to it via its via device; the power
+    # strip in turn has its own child devices (outlets).
+    power_strip_hub = async_create_device(
+        hass, config_entry.entry_id, "Power strip hub", None, None, "power_strip_hub"
+    )
+    power_strip_via_hub = async_create_device(
+        hass,
+        config_entry.entry_id,
+        None,
+        "n_ch_power_strip",
+        {"number_of_sockets": "2"},
+        "2_ch_power_strip_via_hub",
+        via_device_id=power_strip_hub.id,
+    )
+
     async_add_entities(
         [
             DemoSwitch(
@@ -50,6 +68,30 @@ async def async_setup_entry(
                 assumed=False,
                 parent_device_id=parent_device_id,
             ),
+            DemoSwitch(
+                unique_id="hub_switch",
+                device_name="Hub switch",
+                entity_name=None,
+                state=False,
+                assumed=False,
+                via_device_id=hub.id,
+            ),
+            DemoSwitch(
+                unique_id="outlet_3",
+                device_name="Outlet 3",
+                entity_name=None,
+                state=False,
+                assumed=False,
+                parent_device_id=power_strip_via_hub.id,
+            ),
+            DemoSwitch(
+                unique_id="outlet_4",
+                device_name="Outlet 4",
+                entity_name=None,
+                state=True,
+                assumed=False,
+                parent_device_id=power_strip_via_hub.id,
+            ),
         ]
     )
 
@@ -72,6 +114,7 @@ class DemoSwitch(SwitchEntity):
         translation_key: str | None = None,
         device_class: SwitchDeviceClass | None = None,
         parent_device_id: str | None = None,
+        via_device_id: str | None = None,
     ) -> None:
         """Initialize the Demo switch."""
         self._attr_assumed_state = assumed
@@ -90,6 +133,8 @@ class DemoSwitch(SwitchEntity):
                 identifiers={(DOMAIN, unique_id)},
                 name=device_name,
             )
+            if via_device_id is not None:
+                self._attr_device_info["via_device_id"] = via_device_id
         self._attr_name = entity_name
 
     @override

tests/components/kitchen_sink/test_switch.py

@@ -56,6 +56,44 @@ async def test_state(
         assert main_device_entry == snapshot
 
 
+async def test_via_device_relationships(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+) -> None:
+    """Test devices linked to a hub through their via device."""
+    config_entry_id = hass.config_entries.async_entries(DOMAIN)[0].entry_id
+
+    # A standalone switch device linked to a hub through its via device.
+    hub = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "hub"), config_entry_id
+    )
+    assert hub is not None
+    hub_switch = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "hub_switch"), config_entry_id
+    )
+    assert hub_switch is not None
+    assert hub_switch.via_device_id == hub.id
+
+    # A power strip that has child devices and is itself linked to a hub through
+    # its via device.
+    power_strip_hub = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "power_strip_hub"), config_entry_id
+    )
+    assert power_strip_hub is not None
+    power_strip = device_registry.async_get_device_by_identifier(
+        (DOMAIN, "2_ch_power_strip_via_hub"), config_entry_id
+    )
+    assert power_strip is not None
+    assert power_strip.via_device_id == power_strip_hub.id
+
+    for outlet_unique_id in ("outlet_3", "outlet_4"):
+        outlet = device_registry.async_get_child_device_by_identifier(
+            (DOMAIN, outlet_unique_id), config_entry_id
+        )
+        assert outlet is not None
+        assert outlet.parent_device_id == power_strip.id
+
+
 @pytest.mark.parametrize("switch_entity_id", SWITCH_ENTITY_IDS)
 async def test_turn_on(hass: HomeAssistant, switch_entity_id: str) -> None:
     """Test switch turn on method."""