Add wait_for_active_state helper to SystemdUnit

Centralize the repeated pattern of listening for D-Bus
PropertiesChanged signals to wait for a systemd unit's ActiveState
to reach a target state. Refactor core.py, host/firewall.py, and
mounts/mount.py to use the new helper.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

supervisor/core.py

@@ -16,12 +16,7 @@ from .const import (
     CoreState,
 )
 from .coresys import CoreSys, CoreSysAttributes
-from .dbus.const import (
-    DBUS_ATTR_ACTIVE_STATE,
-    DBUS_IFACE_SYSTEMD_UNIT,
-    StopUnitMode,
-    UnitActiveState,
-)
+from .dbus.const import StopUnitMode, UnitActiveState
 from .exceptions import (
     HassioError,
     HomeAssistantCrashError,
@@ -444,25 +439,13 @@ class Core(CoreSysAttributes):
                 "systemd-timesyncd.service"
             )
             try:
-                async with (
-                    asyncio.timeout(10),
-                    timesync_unit.properties_changed() as signal,
-                ):
+                async with asyncio.timeout(10):
                     await self.sys_dbus.systemd.stop_unit(
                         "systemd-timesyncd.service", StopUnitMode.REPLACE
                     )
-                    while (
-                        await timesync_unit.get_active_state()
-                        != UnitActiveState.INACTIVE
-                    ):
-                        prop_change = await signal.wait_for_signal()
-                        if (
-                            prop_change[0] == DBUS_IFACE_SYSTEMD_UNIT
-                            and DBUS_ATTR_ACTIVE_STATE in prop_change[1]
-                            and prop_change[1][DBUS_ATTR_ACTIVE_STATE].value
-                            == UnitActiveState.INACTIVE
-                        ):
-                            break
+                    await timesync_unit.wait_for_active_state(
+                        {UnitActiveState.INACTIVE}
+                    )
             except TimeoutError:
                 _LOGGER.warning(
                     "Timeout waiting for systemd-timesyncd to stop, "

supervisor/dbus/systemd.py

@@ -16,6 +16,7 @@ from ..exceptions import (
 )
 from ..utils.dbus import DBusSignalWrapper
 from .const import (
+    DBUS_ATTR_ACTIVE_STATE,
     DBUS_ATTR_FINISH_TIMESTAMP,
     DBUS_ATTR_FIRMWARE_TIMESTAMP_MONOTONIC,
     DBUS_ATTR_KERNEL_TIMESTAMP_MONOTONIC,
@@ -24,6 +25,7 @@ from .const import (
     DBUS_ATTR_VIRTUALIZATION,
     DBUS_ERR_SYSTEMD_NO_SUCH_UNIT,
     DBUS_IFACE_SYSTEMD_MANAGER,
+    DBUS_IFACE_SYSTEMD_UNIT,
     DBUS_NAME_SYSTEMD,
     DBUS_OBJECT_SYSTEMD,
     DBUS_SIGNAL_PROPERTIES_CHANGED,
@@ -86,6 +88,25 @@ class SystemdUnit(DBusInterface):
         """Return signal wrapper for properties changed."""
         return self.connected_dbus.signal(DBUS_SIGNAL_PROPERTIES_CHANGED)
 
+    @dbus_connected
+    async def wait_for_active_state(
+        self, target_states: set[UnitActiveState]
+    ) -> UnitActiveState:
+        """Wait for unit to reach one of the target active states.
+
+        Caller must handle TimeoutError if a timeout is desired.
+        """
+        async with self.properties_changed() as signal:
+            state = await self.get_active_state()
+            while state not in target_states:
+                interface, changed, _ = await signal.wait_for_signal()
+                if (
+                    interface == DBUS_IFACE_SYSTEMD_UNIT
+                    and DBUS_ATTR_ACTIVE_STATE in changed
+                ):
+                    state = UnitActiveState(changed[DBUS_ATTR_ACTIVE_STATE].value)
+        return state
+
 
 class Systemd(DBusInterfaceProxy):
     """Systemd function handler.

supervisor/host/firewall.py

@@ -8,12 +8,7 @@ from dbus_fast import Variant
 
 from ..const import DOCKER_IPV4_NETWORK_MASK, DOCKER_IPV6_NETWORK_MASK, DOCKER_NETWORK
 from ..coresys import CoreSys, CoreSysAttributes
-from ..dbus.const import (
-    DBUS_ATTR_ACTIVE_STATE,
-    DBUS_IFACE_SYSTEMD_UNIT,
-    StartUnitMode,
-    UnitActiveState,
-)
+from ..dbus.const import StartUnitMode, UnitActiveState
 from ..dbus.systemd import ExecStartEntry
 from ..exceptions import DBusError
 from ..resolution.const import UnhealthyReason
@@ -125,18 +120,8 @@ class FirewallManager(CoreSysAttributes):
         # Wait for the oneshot unit to finish and verify it succeeded
         try:
             unit = await self.sys_dbus.systemd.get_unit(FIREWALL_SERVICE)
-            async with (
-                asyncio.timeout(FIREWALL_UNIT_TIMEOUT),
-                unit.properties_changed() as signal,
-            ):
-                state = await unit.get_active_state()
-                while state not in TERMINAL_STATES:
-                    props = await signal.wait_for_signal()
-                    if (
-                        props[0] == DBUS_IFACE_SYSTEMD_UNIT
-                        and DBUS_ATTR_ACTIVE_STATE in props[1]
-                    ):
-                        state = UnitActiveState(props[1][DBUS_ATTR_ACTIVE_STATE].value)
+            async with asyncio.timeout(FIREWALL_UNIT_TIMEOUT):
+                state = await unit.wait_for_active_state(TERMINAL_STATES)
         except (DBusError, TimeoutError) as err:
             _LOGGER.error(
                 "Failed waiting for gateway firewall unit to complete: %s", err

supervisor/mounts/mount.py

@@ -12,12 +12,10 @@ from voluptuous import Coerce
 
 from ..coresys import CoreSys, CoreSysAttributes
 from ..dbus.const import (
-    DBUS_ATTR_ACTIVE_STATE,
     DBUS_ATTR_DESCRIPTION,
     DBUS_ATTR_OPTIONS,
     DBUS_ATTR_TYPE,
     DBUS_ATTR_WHAT,
-    DBUS_IFACE_SYSTEMD_UNIT,
     StartUnitMode,
     StopUnitMode,
     UnitActiveState,
@@ -179,7 +177,7 @@ class Mount(CoreSysAttributes, ABC):
             await self.mount()
             return
 
-        await self._update_state_await(unit, not_state=UnitActiveState.ACTIVATING)
+        await self._update_state_await(unit)
 
         # If mount is not available, try to reload it
         if not await self.is_mounted():
@@ -225,29 +223,20 @@ class Mount(CoreSysAttributes, ABC):
     async def _update_state_await(
         self,
         unit: SystemdUnit,
-        expected_states: list[UnitActiveState] | None = None,
-        not_state: UnitActiveState = UnitActiveState.ACTIVATING,
+        expected_states: set[UnitActiveState] | None = None,
     ) -> None:
-        """Update state info about mount from dbus. Wait for one of expected_states to appear or state to change from not_state."""
+        """Update state info about mount from dbus. Wait for one of expected_states to appear."""
+        if expected_states is None:
+            expected_states = {
+                UnitActiveState.ACTIVE,
+                UnitActiveState.FAILED,
+                UnitActiveState.INACTIVE,
+            }
         try:
-            async with asyncio.timeout(30), unit.properties_changed() as signal:
-                await self._update_state(unit)
-                while (
-                    expected_states
-                    and self.state not in expected_states
-                    or not expected_states
-                    and self.state == not_state
-                ):
-                    prop_change_signal = await signal.wait_for_signal()
-                    if (
-                        prop_change_signal[0] == DBUS_IFACE_SYSTEMD_UNIT
-                        and DBUS_ATTR_ACTIVE_STATE in prop_change_signal[1]
-                    ):
-                        self._state = prop_change_signal[1][
-                            DBUS_ATTR_ACTIVE_STATE
-                        ].value
-
+            async with asyncio.timeout(30):
+                self._state = await unit.wait_for_active_state(expected_states)
         except TimeoutError:
+            await self._update_state(unit)
             _LOGGER.warning(
                 "Mount %s still in state %s after waiting for 30 seconds to complete",
                 self.name,
@@ -300,7 +289,7 @@ class Mount(CoreSysAttributes, ABC):
             ) from err
 
         if unit := await self._update_unit():
-            await self._update_state_await(unit, not_state=UnitActiveState.ACTIVATING)
+            await self._update_state_await(unit)
 
         if not await self.is_mounted():
             raise MountActivationError(
@@ -320,7 +309,7 @@ class Mount(CoreSysAttributes, ABC):
                 await self.sys_dbus.systemd.stop_unit(self.unit_name, StopUnitMode.FAIL)
 
             await self._update_state_await(
-                unit, [UnitActiveState.INACTIVE, UnitActiveState.FAILED]
+                unit, {UnitActiveState.INACTIVE, UnitActiveState.FAILED}
             )
 
             if self.state == UnitActiveState.FAILED:
@@ -349,9 +338,7 @@ class Mount(CoreSysAttributes, ABC):
             await self._restart()
         else:
             if unit := await self._update_unit():
-                await self._update_state_await(
-                    unit, not_state=UnitActiveState.ACTIVATING
-                )
+                await self._update_state_await(unit)
 
             if not await self.is_mounted():
                 _LOGGER.info(
@@ -380,7 +367,7 @@ class Mount(CoreSysAttributes, ABC):
             ) from err
 
         if unit := await self._update_unit():
-            await self._update_state_await(unit, not_state=UnitActiveState.ACTIVATING)
+            await self._update_state_await(unit)
 
         if not await self.is_mounted():
             raise MountActivationError(