Format data disk in Supervisor instead of OS Agent (#4212)

* Supervisor formats data disk instead of os agent

* Fix issues occurring during tests

* Can't migrate if target is too small

supervisor/dbus/agent/datadisk.py

@@ -33,3 +33,8 @@ class DataDisk(DBusInterfaceProxy):
     async def reload_device(self) -> None:
         """Reload device data."""
         await self.dbus.DataDisk.call_reload_device()
+
+    @dbus_connected
+    async def mark_data_move(self) -> None:
+        """Create marker to signal to do data disk migration next reboot."""
+        await self.dbus.DataDisk.call_mark_data_move()

supervisor/dbus/interface.py

@@ -1,5 +1,6 @@
 """Interface class for D-Bus wrappers."""
 from abc import ABC
+from collections.abc import Callable
 from functools import wraps
 from typing import Any
 
@@ -83,6 +84,7 @@ class DBusInterfaceProxy(DBusInterface):
     properties_interface: str | None = None
     properties: dict[str, Any] | None = None
     sync_properties: bool = True
+    _sync_properties_callback: Callable | None = None
 
     def __init__(self):
         """Initialize properties."""
@@ -104,7 +106,17 @@ class DBusInterfaceProxy(DBusInterface):
 
         await self.update()
         if self.sync_properties and self.is_connected:
-            self.dbus.sync_property_changes(self.properties_interface, self.update)
+            self._sync_properties_callback = self.dbus.sync_property_changes(
+                self.properties_interface, self.update
+            )
+
+    def stop_sync_property_changes(self) -> None:
+        """Stop syncing property changes to object."""
+        if not self._sync_properties_callback:
+            return
+
+        self.dbus.stop_sync_property_changes(self._sync_properties_callback)
+        self.sync_properties = False
 
     @dbus_connected
     async def update(self, changed: dict[str, Any] | None = None) -> None:

supervisor/dbus/udisks2/__init__.py

@@ -139,8 +139,20 @@ class UDisks2(DBusInterfaceProxy):
         resolved = {
             device: self._block_devices[device]
             if device in self._block_devices
+            and self._block_devices[device].is_connected
             else await UDisks2Block.new(device, self.dbus.bus)
             for device in block_devices
         }
         self._block_devices.update(resolved)
         return list(resolved.values())
+
+    def shutdown(self) -> None:
+        """Shutdown the object and disconnect from D-Bus.
+
+        This method is irreversible.
+        """
+        for block_device in self.block_devices:
+            block_device.shutdown()
+        for drive in self.drives:
+            drive.shutdown()
+        super().shutdown()

supervisor/dbus/udisks2/block.py

@@ -1,6 +1,8 @@
 """Interface to UDisks2 Block Device over D-Bus."""
+import asyncio
 from collections.abc import Callable
 from pathlib import Path
+from typing import Any
 
 from awesomeversion import AwesomeVersion
 from dbus_fast.aio import MessageBus
@@ -61,18 +63,7 @@ class UDisks2Block(DBusInterfaceProxy):
     async def connect(self, bus: MessageBus) -> None:
         """Connect to bus."""
         await super().connect(bus)
-
-        if DBUS_IFACE_FILESYSTEM in self.dbus.proxies:
-            self._filesystem = UDisks2Filesystem(
-                self.object_path, sync_properties=self.sync_properties
-            )
-            await self._filesystem.initialize(self.dbus)
-
-        if DBUS_IFACE_PARTITION_TABLE in self.dbus.proxies:
-            self._partition_table = UDisks2PartitionTable(
-                self.object_path, sync_properties=self.sync_properties
-            )
-            await self._partition_table.initialize(self.dbus)
+        await self._reload_interfaces()
 
     @staticmethod
     async def new(
@@ -195,6 +186,58 @@ class UDisks2Block(DBusInterfaceProxy):
         """
         return self.properties[DBUS_ATTR_DRIVE]
 
+    @dbus_connected
+    async def update(self, changed: dict[str, Any] | None = None) -> None:
+        """Update properties via D-Bus."""
+        await super().update(changed)
+
+        if not changed:
+            await asyncio.gather(
+                *[
+                    intr.update()
+                    for intr in (self.filesystem, self.partition_table)
+                    if intr
+                ]
+            )
+
+    @dbus_connected
+    async def check_type(self) -> None:
+        """Check if type of block device has changed and adjust interfaces if so."""
+        introspection = await self.dbus.introspect()
+        interfaces = {intr.name for intr in introspection.interfaces}
+
+        # If interfaces changed, update the proxy from introspection and reload interfaces
+        if interfaces != set(self.dbus.proxies.keys()):
+            await self.dbus.init_proxy(introspection=introspection)
+            await self._reload_interfaces()
+
+    @dbus_connected
+    async def _reload_interfaces(self) -> None:
+        """Reload interfaces from introspection as necessary."""
+        # Check if block device is a filesystem
+        if not self.filesystem and DBUS_IFACE_FILESYSTEM in self.dbus.proxies:
+            self._filesystem = UDisks2Filesystem(
+                self.object_path, sync_properties=self.sync_properties
+            )
+            await self._filesystem.initialize(self.dbus)
+
+        elif self.filesystem and DBUS_IFACE_FILESYSTEM not in self.dbus.proxies:
+            self.filesystem.stop_sync_property_changes()
+            self._filesystem = None
+
+        # Check if block device is a partition table
+        if not self.partition_table and DBUS_IFACE_PARTITION_TABLE in self.dbus.proxies:
+            self._partition_table = UDisks2PartitionTable(
+                self.object_path, sync_properties=self.sync_properties
+            )
+            await self._partition_table.initialize(self.dbus)
+
+        elif (
+            self.partition_table and DBUS_IFACE_PARTITION_TABLE not in self.dbus.proxies
+        ):
+            self.partition_table.stop_sync_property_changes()
+            self._partition_table = None
+
     @dbus_connected
     async def format(
         self, type_: FormatType = FormatType.GPT, options: FormatOptions | None = None

supervisor/dbus/udisks2/data.py

@@ -82,7 +82,7 @@ class DeviceSpecification:
     def to_dict(self) -> dict[str, Variant]:
         """Return dict representation."""
         data = {
-            "path": Variant("s", str(self.path)) if self.path else None,
+            "path": Variant("s", self.path.as_posix()) if self.path else None,
             "label": _optional_variant("s", self.label),
             "uuid": _optional_variant("s", self.uuid),
         }

supervisor/dbus/udisks2/partition_table.py

@@ -48,14 +48,15 @@ class UDisks2PartitionTable(DBusInterfaceProxy):
         self,
         offset: int = 0,
         size: int = 0,
-        type_: PartitionTableType = PartitionTableType.UNKNOWN,
+        type_: str = "",
         name: str = "",
         options: CreatePartitionOptions | None = None,
     ) -> str:
         """Create a new partition and return object path of new block device.
 
-        'UNKNOWN' for type here means UDisks2 selects default based on partition table and OS.
-        Use with UDisks2Block.new to get block object. Or UDisks2.get_block_device after UDisks2.update.
+        Type should be a GUID from https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_type_GUIDs
+        or empty string and let UDisks2 choose a default based on partition table and OS.
+        Provide return value with UDisks2Block.new. Or UDisks2.get_block_device after UDisks2.update.
         """
         options = options.to_dict() if options else {}
         return await self.dbus.PartitionTable.call_create_partition(

supervisor/os/data_disk.py

@@ -4,11 +4,13 @@ from contextlib import suppress
 from dataclasses import dataclass
 import logging
 from pathlib import Path
+from typing import Final
 
 from awesomeversion import AwesomeVersion
 
 from ..coresys import CoreSys, CoreSysAttributes
 from ..dbus.udisks2.block import UDisks2Block
+from ..dbus.udisks2.const import FormatType
 from ..dbus.udisks2.drive import UDisks2Drive
 from ..exceptions import (
     DBusError,
@@ -20,6 +22,11 @@ from ..exceptions import (
 )
 from ..jobs.const import JobCondition, JobExecutionLimit
 from ..jobs.decorator import Job
+from ..utils.sentry import capture_exception
+
+LINUX_DATA_PARTITION_GUID: Final = "0FC63DAF-8483-4772-8E79-3D69D8477DE4"
+EXTERNAL_DATA_DISK_PARTITION_NAME: Final = "hassos-data-external"
+OS_AGENT_MARK_DATA_MOVE_VERSION: Final = AwesomeVersion("1.5.0")
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -35,6 +42,7 @@ class Disk:
     size: int
     device_path: Path
     object_path: str
+    device_object_path: str
 
     @staticmethod
     def from_udisks2_drive(
@@ -49,6 +57,7 @@ class Disk:
             size=drive.size,
             device_path=drive_block_device.device,
             object_path=drive.object_path,
+            device_object_path=drive_block_device.object_path,
         )
 
     @property
@@ -115,6 +124,7 @@ class DataDisk(CoreSysAttributes):
             size=0,
             device_path=self.sys_dbus.agent.datadisk.current_device,
             object_path="",
+            device_object_path="",
         )
 
     @property
@@ -163,24 +173,42 @@ class DataDisk(CoreSysAttributes):
         # Force a dbus update first so all info is up to date
         await self.sys_dbus.udisks2.update()
 
-        try:
-            target_disk: Disk = next(
-                disk
-                for disk in self.available_disks
-                if disk.id == new_disk or disk.device_path.as_posix() == new_disk
-            )
-        except StopIteration:
+        target_disk: list[Disk] = [
+            disk
+            for disk in self.available_disks
+            if disk.id == new_disk or disk.device_path.as_posix() == new_disk
+        ]
+        if len(target_disk) != 1:
             raise HassOSDataDiskError(
                 f"'{new_disk!s}' not a valid data disk target!", _LOGGER.error
             ) from None
 
-        # Migrate data on Host
-        try:
-            await self.sys_dbus.agent.datadisk.change_device(target_disk.device_path)
-        except DBusError as err:
-            raise HassOSDataDiskError(
-                f"Can't move data partition to {new_disk!s}: {err!s}", _LOGGER.error
-            ) from err
+        # Older OS did not have mark data move API. Must let OS do disk format & migration
+        if self.sys_dbus.agent.version < OS_AGENT_MARK_DATA_MOVE_VERSION:
+            try:
+                await self.sys_dbus.agent.datadisk.change_device(
+                    target_disk[0].device_path
+                )
+            except DBusError as err:
+                raise HassOSDataDiskError(
+                    f"Can't move data partition to {new_disk!s}: {err!s}", _LOGGER.error
+                ) from err
+        else:
+            # Format disk then tell OS to migrate next reboot
+            partition = await self._format_device_with_single_partition(target_disk[0])
+
+            if self.disk_used and partition.size < self.disk_used.size:
+                raise HassOSDataDiskError(
+                    f"Cannot use {new_disk} as data disk as it is smaller then the current one (new: {partition.size}, current: {self.disk_used.size})"
+                )
+
+            try:
+                await self.sys_dbus.agent.datadisk.mark_data_move()
+            except DBusError as err:
+                raise HassOSDataDiskError(
+                    f"Unable to create data disk migration marker: {err!s}",
+                    _LOGGER.error,
+                ) from err
 
         # Restart Host for finish the process
         try:
@@ -190,3 +218,50 @@ class DataDisk(CoreSysAttributes):
                 f"Can't restart device to finish disk migration: {err!s}",
                 _LOGGER.warning,
             ) from err
+
+    async def _format_device_with_single_partition(
+        self, new_disk: Disk
+    ) -> UDisks2Block:
+        """Format device with a single partition to use as data disk."""
+        block_device: UDisks2Block = self.sys_dbus.udisks2.get_block_device(
+            new_disk.device_object_path
+        )
+
+        try:
+            await block_device.format(FormatType.GPT)
+        except DBusError as err:
+            capture_exception(err)
+            raise HassOSDataDiskError(
+                f"Could not format {new_disk.id}: {err!s}", _LOGGER.error
+            ) from err
+
+        await block_device.check_type()
+        if not block_device.partition_table:
+            raise HassOSDataDiskError(
+                "Block device does not contain a partition table after format, cannot create data partition",
+                _LOGGER.error,
+            )
+
+        try:
+            partition = await block_device.partition_table.create_partition(
+                0, 0, LINUX_DATA_PARTITION_GUID, EXTERNAL_DATA_DISK_PARTITION_NAME
+            )
+        except DBusError as err:
+            capture_exception(err)
+            raise HassOSDataDiskError(
+                f"Could not create new data partition: {err!s}"
+            ) from err
+
+        try:
+            partition_block = await UDisks2Block.new(
+                partition, self.sys_dbus.bus, sync_properties=False
+            )
+        except DBusError as err:
+            raise HassOSDataDiskError(
+                f"New data partition at {partition} is missing or unusable"
+            ) from err
+
+        _LOGGER.debug(
+            "New data partition prepared on device %s", partition_block.device
+        )
+        return partition_block

supervisor/utils/dbus.py

@@ -56,7 +56,7 @@ class DBus:
         self = DBus(bus, bus_name, object_path)
 
         # pylint: disable=protected-access
-        await self._init_proxy()
+        await self.init_proxy()
 
         _LOGGER.debug("Connect to D-Bus: %s - %s", bus_name, object_path)
         return self
@@ -115,13 +115,11 @@ class DBus:
             for interface in self._proxy_obj.introspection.interfaces
         }
 
-    async def _init_proxy(self) -> None:
-        """Read interface data."""
-        introspection: Node | None = None
-
+    async def introspect(self) -> Node:
+        """Return introspection for dbus object."""
         for _ in range(3):
             try:
-                introspection = await self._bus.introspect(
+                return await self._bus.introspect(
                     self.bus_name, self.object_path, timeout=10
                 )
             except InvalidIntrospectionError as err:
@@ -132,21 +130,41 @@ class DBus:
                 _LOGGER.warning(
                     "Busy system at %s - %s", self.bus_name, self.object_path
                 )
-            else:
-                break
 
             await asyncio.sleep(3)
 
-        if introspection is None:
-            raise DBusFatalError(
-                "Could not get introspection data after 3 attempts", _LOGGER.error
-            )
+        raise DBusFatalError(
+            "Could not get introspection data after 3 attempts", _LOGGER.error
+        )
+
+    async def init_proxy(self, *, introspection: Node | None = None) -> None:
+        """Read interface data."""
+        if not introspection:
+            introspection = await self.introspect()
+
+        # If we have a proxy obj store signal monitors and disconnect first
+        signal_monitors = self._signal_monitors
+        if self._proxy_obj:
+            self.disconnect()
 
         self._proxy_obj = self.bus.get_proxy_object(
             self.bus_name, self.object_path, introspection
         )
         self._add_interfaces()
 
+        # Reconnect existing signal monitors on new proxy obj if possible (introspection may have changed)
+        for intr, signals in signal_monitors.items():
+            for name, callbacks in signals.items():
+                if intr in self._proxies and hasattr(self._proxies[intr], f"on_{name}"):
+                    for callback in callbacks:
+                        try:
+                            getattr(self._proxies[intr], f"on_{name}")(
+                                callback, unpack_variants=True
+                            )
+                            self._add_signal_monitor(intr, name, callback)
+                        except Exception:  # pylint: disable=broad-except
+                            _LOGGER.exception("Can't re-add signal listener")
+
     @property
     def proxies(self) -> dict[str, ProxyInterface]:
         """Return all proxies."""
@@ -232,6 +250,19 @@ class DBus:
         """Get signal context manager for this object."""
         return DBusSignalWrapper(self, signal_member)
 
+    def _add_signal_monitor(
+        self, interface: str, dbus_name: str, callback: Callable
+    ) -> None:
+        """Add a callback to the tracked signal monitors."""
+        if interface not in self._signal_monitors:
+            self._signal_monitors[interface] = {}
+
+        if dbus_name not in self._signal_monitors[interface]:
+            self._signal_monitors[interface][dbus_name] = [callback]
+
+        else:
+            self._signal_monitors[interface][dbus_name].append(callback)
+
     def __getattr__(self, name: str) -> DBusCallWrapper:
         """Map to dbus method."""
         return getattr(DBusCallWrapper(self, self.bus_name), name)
@@ -286,17 +317,8 @@ class DBusCallWrapper:
                     getattr(self._proxy, name)(callback, unpack_variants=True)
 
                     # pylint: disable=protected-access
-                    if self.interface not in self.dbus._signal_monitors:
-                        self.dbus._signal_monitors[self.interface] = {}
-
-                    if dbus_name not in self.dbus._signal_monitors[self.interface]:
-                        self.dbus._signal_monitors[self.interface][dbus_name] = [
-                            callback
-                        ]
-                    else:
-                        self.dbus._signal_monitors[self.interface][dbus_name].append(
-                            callback
-                        )
+                    self.dbus._add_signal_monitor(self.interface, dbus_name, callback)
+                    # pylint: enable=protected-access
 
                 return _on_signal
 
@@ -322,6 +344,7 @@ class DBusCallWrapper:
                         del self.dbus._signal_monitors[self.interface][dbus_name]
                         if not self.dbus._signal_monitors[self.interface]:
                             del self.dbus._signal_monitors[self.interface]
+                # pylint: enable=protected-access
 
             return _off_signal
 

tests/conftest.py

@@ -310,8 +310,6 @@ async def coresys(
     ):
         message_bus.return_value.connect = AsyncMock(return_value=dbus_session_bus)
         await coresys_obj._dbus.load()
-    # coresys_obj._dbus._bus = dbus_session_bus
-    # coresys_obj._dbus._network = network_manager
 
     # Mock docker
     coresys_obj._docker = docker
@@ -350,6 +348,7 @@ async def coresys(
     with patch("supervisor.store.git.GitRepo._remove"):
         yield coresys_obj
 
+    await coresys_obj.dbus.unload()
     await coresys_obj.websession.close()
 
 

tests/dbus/agent/test_datadisk.py

@@ -72,3 +72,19 @@ async def test_dbus_osagent_datadisk_reload_device(
 
     assert await os_agent.datadisk.reload_device() is None
     assert datadisk_service.ReloadDevice.calls == [tuple()]
+
+
+async def test_dbus_osagent_datadisk_mark_data_move(
+    datadisk_service: DataDiskService, dbus_session_bus: MessageBus
+):
+    """Create data disk migration marker for next reboot."""
+    datadisk_service.MarkDataMove.calls.clear()
+    os_agent = OSAgent()
+
+    with pytest.raises(DBusNotConnectedError):
+        await os_agent.datadisk.mark_data_move()
+
+    await os_agent.connect(dbus_session_bus)
+
+    assert await os_agent.datadisk.mark_data_move() is None
+    assert datadisk_service.MarkDataMove.calls == [tuple()]

tests/dbus/network/test_network_manager.py

@@ -142,7 +142,7 @@ async def test_handling_bad_devices(
     caplog.clear()
     caplog.set_level(logging.INFO, "supervisor.dbus.network")
 
-    with patch.object(DBus, "_init_proxy", side_effect=DBusFatalError()):
+    with patch.object(DBus, "init_proxy", side_effect=DBusFatalError()):
         await network_manager.update(
             {"Devices": ["/org/freedesktop/NetworkManager/Devices/100"]}
         )
@@ -157,7 +157,7 @@ async def test_handling_bad_devices(
 
     # Unparseable introspections shouldn't happen, this one is logged and captured
     await network_manager.update()
-    with patch.object(DBus, "_init_proxy", side_effect=(err := DBusParseError())):
+    with patch.object(DBus, "init_proxy", side_effect=(err := DBusParseError())):
         await network_manager.update(
             {"Devices": [device := "/org/freedesktop/NetworkManager/Devices/102"]}
         )
@@ -167,7 +167,7 @@ async def test_handling_bad_devices(
     # We should be able to debug these situations if necessary
     caplog.set_level(logging.DEBUG, "supervisor.dbus.network")
     await network_manager.update()
-    with patch.object(DBus, "_init_proxy", side_effect=DBusFatalError()):
+    with patch.object(DBus, "init_proxy", side_effect=DBusFatalError()):
         await network_manager.update(
             {"Devices": [device := "/org/freedesktop/NetworkManager/Devices/103"]}
         )

tests/dbus/test_interface.py

@@ -56,13 +56,12 @@ async def fixture_proxy(
     proxy.bus_name = "service.test.TestInterface"
     proxy.object_path = "/service/test/TestInterface"
     proxy.properties_interface = "service.test.TestInterface"
-    proxy.sync_properties = request.param
+    proxy.sync_properties = getattr(request, "param", True)
 
     await proxy.connect(dbus_session_bus)
     yield proxy
 
 
-@pytest.mark.parametrize("proxy", [True], indirect=True)
 async def test_dbus_proxy_connect(
     proxy: DBusInterfaceProxy, test_service: TestInterface
 ):
@@ -213,3 +212,21 @@ async def test_initialize(test_service: TestInterface, dbus_session_bus: Message
         )
     )
     assert proxy.is_connected is True
+
+
+async def test_stop_sync_property_changes(
+    proxy: DBusInterfaceProxy, test_service: TestInterface
+):
+    """Test stop sync property changes disables the sync via signal."""
+    assert proxy.is_connected
+    assert proxy.properties["TestProp"] == 4
+
+    test_service.emit_properties_changed({"TestProp": 1})
+    await test_service.ping()
+    assert proxy.properties["TestProp"] == 1
+
+    proxy.stop_sync_property_changes()
+
+    test_service.emit_properties_changed({"TestProp": 4})
+    await test_service.ping()
+    assert proxy.properties["TestProp"] == 1

tests/dbus/udisks2/test_block.py

@@ -1,6 +1,7 @@
 """Test UDisks2 Block Device interface."""
 
 from pathlib import Path
+from unittest.mock import patch
 
 from dbus_fast import Variant
 from dbus_fast.aio.message_bus import MessageBus
@@ -9,7 +10,11 @@ import pytest
 from supervisor.dbus.udisks2.block import UDisks2Block
 from supervisor.dbus.udisks2.const import FormatType, PartitionTableType
 from supervisor.dbus.udisks2.data import FormatOptions
+from supervisor.dbus.udisks2.filesystem import UDisks2Filesystem
+from supervisor.dbus.udisks2.partition_table import UDisks2PartitionTable
+from supervisor.utils.dbus import DBus
 
+from tests.common import mock_dbus_services
 from tests.dbus_service_mocks.base import DBusServiceMock
 from tests.dbus_service_mocks.udisks2_block import Block as BlockService
 
@@ -107,3 +112,92 @@ async def test_format(block_sda_service: BlockService, dbus_session_bus: Message
             },
         )
     ]
+
+
+async def test_check_type(dbus_session_bus: MessageBus):
+    """Test block device changes types correctly."""
+    block_services = (
+        await mock_dbus_services(
+            {
+                "udisks2_block": [
+                    "/org/freedesktop/UDisks2/block_devices/sda",
+                    "/org/freedesktop/UDisks2/block_devices/sda1",
+                ]
+            },
+            dbus_session_bus,
+        )
+    )["udisks2_block"]
+    sda_block_service = block_services["/org/freedesktop/UDisks2/block_devices/sda"]
+    sda1_block_service = block_services["/org/freedesktop/UDisks2/block_devices/sda1"]
+
+    sda = UDisks2Block("/org/freedesktop/UDisks2/block_devices/sda")
+    sda1 = UDisks2Block("/org/freedesktop/UDisks2/block_devices/sda1")
+    await sda.connect(dbus_session_bus)
+    await sda1.connect(dbus_session_bus)
+
+    # Connected but neither are filesystems are partition tables
+    assert sda.partition_table is None
+    assert sda1.filesystem is None
+    assert sda.id_label == ""
+    assert sda1.id_label == "hassos-data"
+
+    # Store current introspection then make sda into a partition table and sda1 into a filesystem
+    orig_introspection = await sda.dbus.introspect()
+    services = await mock_dbus_services(
+        {
+            "udisks2_partition_table": "/org/freedesktop/UDisks2/block_devices/sda",
+            "udisks2_filesystem": "/org/freedesktop/UDisks2/block_devices/sda1",
+        },
+        dbus_session_bus,
+    )
+    sda_pt_service = services["udisks2_partition_table"]
+    sda1_fs_service = services["udisks2_filesystem"]
+
+    await sda.check_type()
+    await sda1.check_type()
+
+    # Check that the type is now correct and property changes are syncing
+    assert sda.partition_table
+    assert sda1.filesystem
+
+    partition_table: UDisks2PartitionTable = sda.partition_table
+    filesystem: UDisks2Filesystem = sda1.filesystem
+    assert partition_table.type == PartitionTableType.GPT
+    assert filesystem.size == 250058113024
+
+    sda_pt_service.emit_properties_changed({"Type": "dos"})
+    await sda_pt_service.ping()
+    assert partition_table.type == PartitionTableType.DOS
+
+    sda1_fs_service.emit_properties_changed({"Size": 100})
+    await sda1_fs_service.ping()
+    assert filesystem.size == 100
+
+    # Force introspection to return the original block device only introspection and re-check type
+    with patch.object(DBus, "introspect", return_value=orig_introspection):
+        await sda.check_type()
+        await sda1.check_type()
+
+    # Check that it's a connected block device and no longer the other types
+    assert sda.is_connected is True
+    assert sda1.is_connected is True
+    assert sda.partition_table is None
+    assert sda1.filesystem is None
+
+    # Property changes should still sync for the block devices
+    sda_block_service.emit_properties_changed({"IdLabel": "test"})
+    await sda_block_service.ping()
+    assert sda.id_label == "test"
+
+    sda1_block_service.emit_properties_changed({"IdLabel": "test"})
+    await sda1_block_service.ping()
+    assert sda1.id_label == "test"
+
+    # Property changes should stop syncing for the now unused dbus objects
+    sda_pt_service.emit_properties_changed({"Type": "gpt"})
+    await sda_pt_service.ping()
+    assert partition_table.type == PartitionTableType.DOS
+
+    sda1_fs_service.emit_properties_changed({"Size": 250058113024})
+    await sda1_fs_service.ping()
+    assert filesystem.size == 100

tests/dbus/udisks2/test_manager.py

@@ -1,5 +1,7 @@
 """Test UDisks2 Manager interface."""
 
+from pathlib import Path
+
 from awesomeversion import AwesomeVersion
 from dbus_fast import Variant
 from dbus_fast.aio.message_bus import MessageBus
@@ -72,6 +74,7 @@ async def test_udisks2_manager_info(
     udisks2_manager_service.emit_properties_changed({}, ["SupportedFilesystems"])
     await udisks2_manager_service.ping()
     await udisks2_manager_service.ping()
+    await udisks2_manager_service.ping()  # Three pings: signal, get all properties and get block devices
     assert udisks2.supported_filesystems == [
         "ext4",
         "vfat",
@@ -126,11 +129,11 @@ async def test_resolve_device(
     udisks2 = UDisks2()
 
     with pytest.raises(DBusNotConnectedError):
-        await udisks2.resolve_device(DeviceSpecification(path="/dev/sda1"))
+        await udisks2.resolve_device(DeviceSpecification(path=Path("/dev/sda1")))
 
     await udisks2.connect(dbus_session_bus)
 
-    devices = await udisks2.resolve_device(DeviceSpecification(path="/dev/sda1"))
+    devices = await udisks2.resolve_device(DeviceSpecification(path=Path("/dev/sda1")))
     assert len(devices) == 1
     assert devices[0].id_label == "hassos-data"
     assert udisks2_manager_service.ResolveDevice.calls == [

tests/dbus/udisks2/test_partition_table.py

@@ -108,17 +108,17 @@ async def test_create_partition(
         await sda.create_partition(
             offset=0,
             size=1000000,
-            type_=PartitionTableType.DOS,
+            type_="0FC63DAF-8483-4772-8E79-3D69D8477DE4",
             name="hassos-data",
             options=CreatePartitionOptions(partition_type="primary"),
         )
-        == "/org/freedesktop/UDisks2/block_devices/sda2"
+        == "/org/freedesktop/UDisks2/block_devices/sda1"
     )
     assert partition_table_sda_service.CreatePartition.calls == [
         (
             0,
             1000000,
-            "dos",
+            "0FC63DAF-8483-4772-8E79-3D69D8477DE4",
             "hassos-data",
             {
                 "partition-type": Variant("s", "primary"),

tests/dbus_service_mocks/agent_datadisk.py

@@ -38,3 +38,7 @@ class DataDisk(DBusServiceMock):
     def ReloadDevice(self) -> "b":
         """Reload device."""
         return True
+
+    @dbus_method()
+    def MarkDataMove(self) -> None:
+        """Do MarkDataMove method."""

tests/dbus_service_mocks/udisks2_partition_table.py

@@ -57,6 +57,7 @@ class PartitionTable(DBusServiceMock):
     """
 
     interface = "org.freedesktop.UDisks2.PartitionTable"
+    new_partition = "/org/freedesktop/UDisks2/block_devices/sda1"
 
     def __init__(self, object_path: str):
         """Initialize object."""
@@ -79,7 +80,7 @@ class PartitionTable(DBusServiceMock):
         self, offset: "t", size: "t", type_: "s", name: "s", options: "a{sv}"
     ) -> "o":
         """Do CreatePartition method."""
-        return "/org/freedesktop/UDisks2/block_devices/sda2"
+        return self.new_partition
 
     @dbus_method()
     def CreatePartitionAndFormat(
@@ -93,4 +94,4 @@ class PartitionTable(DBusServiceMock):
         format_options: "a{sv}",
     ) -> "o":
         """Do CreatePartitionAndFormat method."""
-        return "/org/freedesktop/UDisks2/block_devices/sda2"
+        return self.new_partition

tests/os/test_data_disk.py

@@ -2,6 +2,7 @@
 from pathlib import PosixPath
 from unittest.mock import patch
 
+from dbus_fast import Variant
 import pytest
 
 from supervisor.core import Core
@@ -13,6 +14,10 @@ from tests.common import mock_dbus_services
 from tests.dbus_service_mocks.agent_datadisk import DataDisk as DataDiskService
 from tests.dbus_service_mocks.base import DBusServiceMock
 from tests.dbus_service_mocks.logind import Logind as LogindService
+from tests.dbus_service_mocks.udisks2_block import Block as BlockService
+from tests.dbus_service_mocks.udisks2_partition_table import (
+    PartitionTable as PartitionTableService,
+)
 
 # pylint: disable=protected-access
 
@@ -21,7 +26,7 @@ from tests.dbus_service_mocks.logind import Logind as LogindService
 async def add_unusable_drive(
     coresys: CoreSys,
     udisks2_services: dict[str, DBusServiceMock | dict[str, DBusServiceMock]],
-):
+) -> None:
     """Add mock drive with multiple partition tables for negative tests."""
     await mock_dbus_services(
         {
@@ -53,6 +58,7 @@ async def tests_datadisk_current(coresys: CoreSys):
         size=31268536320,
         device_path=PosixPath("/dev/mmcblk1"),
         object_path="/org/freedesktop/UDisks2/drives/BJTD4R_0x97cde291",
+        device_object_path="/org/freedesktop/UDisks2/block_devices/mmcblk1",
     )
 
 
@@ -87,6 +93,7 @@ async def test_datadisk_list(coresys: CoreSys):
             size=250059350016,
             device_path=PosixPath("/dev/sda"),
             object_path="/org/freedesktop/UDisks2/drives/SSK_SSK_Storage_DF56419883D56",
+            device_object_path="/org/freedesktop/UDisks2/block_devices/sda",
         )
     ]
 
@@ -114,3 +121,83 @@ async def test_datadisk_migrate(
 
     assert datadisk_service.ChangeDevice.calls == [("/dev/sda",)]
     assert logind_service.Reboot.calls == [(False,)]
+
+
+@pytest.mark.parametrize(
+    "new_disk",
+    ["SSK-SSK-Storage-DF56419883D56", "/dev/sda"],
+    ids=["by drive id", "by device path"],
+)
+async def test_datadisk_migrate_mark_data_move(
+    coresys: CoreSys,
+    all_dbus_services: dict[str, DBusServiceMock | dict[str, DBusServiceMock]],
+    new_disk: str,
+):
+    """Test migrating data disk with os agent 1.5.0 or later."""
+    datadisk_service: DataDiskService = all_dbus_services["agent_datadisk"]
+    datadisk_service.ChangeDevice.calls.clear()
+    datadisk_service.MarkDataMove.calls.clear()
+    block_service: BlockService = all_dbus_services["udisks2_block"][
+        "/org/freedesktop/UDisks2/block_devices/sda"
+    ]
+    block_service.Format.calls.clear()
+    partition_table_service: PartitionTableService = all_dbus_services[
+        "udisks2_partition_table"
+    ]["/org/freedesktop/UDisks2/block_devices/sda"]
+    partition_table_service.CreatePartition.calls.clear()
+    logind_service: LogindService = all_dbus_services["logind"]
+    logind_service.Reboot.calls.clear()
+
+    all_dbus_services["os_agent"].emit_properties_changed({"Version": "1.5.0"})
+    await all_dbus_services["os_agent"].ping()
+    coresys.os._available = True
+
+    with patch.object(Core, "shutdown") as shutdown:
+        await coresys.os.datadisk.migrate_disk(new_disk)
+        shutdown.assert_called_once()
+
+    assert datadisk_service.ChangeDevice.calls == []
+    assert datadisk_service.MarkDataMove.calls == [tuple()]
+    assert block_service.Format.calls == [
+        ("gpt", {"auth.no_user_interaction": Variant("b", True)})
+    ]
+    assert partition_table_service.CreatePartition.calls == [
+        (
+            0,
+            0,
+            "0FC63DAF-8483-4772-8E79-3D69D8477DE4",
+            "hassos-data-external",
+            {"auth.no_user_interaction": Variant("b", True)},
+        )
+    ]
+    assert logind_service.Reboot.calls == [(False,)]
+
+
+async def test_datadisk_migrate_too_small(
+    coresys: CoreSys,
+    all_dbus_services: dict[str, DBusServiceMock | dict[str, DBusServiceMock]],
+):
+    """Test migration stops and exits if new partition is too small."""
+    datadisk_service: DataDiskService = all_dbus_services["agent_datadisk"]
+    datadisk_service.MarkDataMove.calls.clear()
+    logind_service: LogindService = all_dbus_services["logind"]
+    logind_service.Reboot.calls.clear()
+
+    partition_table_service: PartitionTableService = all_dbus_services[
+        "udisks2_partition_table"
+    ]["/org/freedesktop/UDisks2/block_devices/sda"]
+    partition_table_service.CreatePartition.calls.clear()
+    partition_table_service.new_partition = (
+        "/org/freedesktop/UDisks2/block_devices/mmcblk1p3"
+    )
+
+    all_dbus_services["os_agent"].emit_properties_changed({"Version": "1.5.0"})
+    await all_dbus_services["os_agent"].ping()
+    coresys.os._available = True
+
+    with pytest.raises(HassOSDataDiskError):
+        await coresys.os.datadisk.migrate_disk("SSK-SSK-Storage-DF56419883D56")
+
+    assert partition_table_service.CreatePartition.calls
+    assert datadisk_service.MarkDataMove.calls == []
+    assert logind_service.Reboot.calls == []

tests/utils/test_dbus.py

@@ -3,7 +3,7 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 from dbus_fast.aio.message_bus import MessageBus
-from dbus_fast.service import method
+from dbus_fast.service import method, signal
 import pytest
 
 from supervisor.dbus.const import DBUS_OBJECT_BASE
@@ -24,6 +24,10 @@ class TestInterface(DBusServiceMock):
     def test(self, _: "b") -> None:  # noqa: F821
         """Do Test method."""
 
+    @signal(name="Test")
+    def signal_test(self) -> None:
+        """Signal Test."""
+
 
 @pytest.fixture(name="test_service")
 async def fixture_test_service(dbus_session_bus: MessageBus) -> TestInterface:
@@ -74,3 +78,96 @@ async def test_internal_dbus_errors(
         await test_obj.call_test(True)
 
     capture_exception.assert_called_once_with(err)
+
+
+async def test_introspect(test_service: TestInterface, dbus_session_bus: MessageBus):
+    """Test introspect of dbus object."""
+    test_obj = DBus(dbus_session_bus, "service.test.TestInterface", DBUS_OBJECT_BASE)
+
+    introspection = await test_obj.introspect()
+
+    assert {"service.test.TestInterface", "org.freedesktop.DBus.Properties"} <= {
+        interface.name for interface in introspection.interfaces
+    }
+    test_interface = next(
+        interface
+        for interface in introspection.interfaces
+        if interface.name == "service.test.TestInterface"
+    )
+    assert "Test" in {method_.name for method_ in test_interface.methods}
+
+
+async def test_init_proxy(test_service: TestInterface, dbus_session_bus: MessageBus):
+    """Test init proxy on already connected object to update interfaces."""
+    test_obj = await DBus.connect(
+        dbus_session_bus, "service.test.TestInterface", DBUS_OBJECT_BASE
+    )
+    orig_introspection = await test_obj.introspect()
+    callback_count = 0
+
+    def test_callback():
+        nonlocal callback_count
+        callback_count += 1
+
+    class TestInterface2(TestInterface):
+        """Test interface 2."""
+
+        interface = "service.test.TestInterface.Test2"
+        object_path = DBUS_OBJECT_BASE
+
+    # Test interfaces and methods match expected
+    assert "service.test.TestInterface" in test_obj.proxies
+    assert await test_obj.call_test(True) is None
+    assert "service.test.TestInterface.Test2" not in test_obj.proxies
+
+    # Test basic signal listening works
+    test_obj.on_test(test_callback)
+    test_service.signal_test()
+    await test_service.ping()
+    assert callback_count == 1
+    callback_count = 0
+
+    # Export the second interface and re-create proxy
+    test_service_2 = TestInterface2()
+    test_service_2.export(dbus_session_bus)
+
+    await test_obj.init_proxy()
+
+    # Test interfaces and methods match expected
+    assert "service.test.TestInterface" in test_obj.proxies
+    assert await test_obj.call_test(True) is None
+    assert "service.test.TestInterface.Test2" in test_obj.proxies
+    assert await test_obj.Test2.call_test(True) is None
+
+    # Test signal listening. First listener should still be attached
+    test_obj.Test2.on_test(test_callback)
+    test_service_2.signal_test()
+    await test_service_2.ping()
+    assert callback_count == 1
+
+    test_service.signal_test()
+    await test_service.ping()
+    assert callback_count == 2
+    callback_count = 0
+
+    # Return to original introspection and test interfaces have reset
+    await test_obj.init_proxy(introspection=orig_introspection)
+
+    assert "service.test.TestInterface" in test_obj.proxies
+    assert "service.test.TestInterface.Test2" not in test_obj.proxies
+
+    # Signal listener for second interface should disconnect, first remains
+    test_service_2.signal_test()
+    await test_service_2.ping()
+    assert callback_count == 0
+
+    test_service.signal_test()
+    await test_service.ping()
+    assert callback_count == 1
+    callback_count = 0
+
+    # Should be able to disconnect first signal listener on new proxy obj
+    test_obj.off_test(test_callback)
+    test_service.signal_test()
+    await test_service.ping()
+    assert callback_count == 0