diff --git a/subiquity/server/controllers/filesystem.py b/subiquity/server/controllers/filesystem.py index f210267c..fbdaeebe 100644 --- a/subiquity/server/controllers/filesystem.py +++ b/subiquity/server/controllers/filesystem.py @@ -275,7 +275,7 @@ class FilesystemController(SubiquityController, FilesystemManipulator): self._on_volume: Optional[snapdapi.OnVolume] = None self._source_handler: Optional[AbstractSourceHandler] = None self._system_mounter: Optional[Mounter] = None - self._role_to_device: Dict[str, _Device] = {} + self._role_to_device: Dict[Union[str, snapdapi.Role], _Device] = {} self._device_to_structure: Dict[_Device, snapdapi.OnVolume] = {} self._pyudev_context: Optional[pyudev.Context] = None self.use_tpm: bool = False @@ -949,9 +949,9 @@ class FilesystemController(SubiquityController, FilesystemManipulator): step=snapdapi.SystemActionStep.SETUP_STORAGE_ENCRYPTION, on_volumes=self._on_volumes(), ), + ann=snapdapi.SystemActionResponse, ) - role_to_encrypted_device = result["encrypted-devices"] - for role, enc_path in role_to_encrypted_device.items(): + for role, enc_path in result.encrypted_devices.items(): arb_device = ArbitraryDevice(m=self.model, path=enc_path) self.model._actions.append(arb_device) part = self._role_to_device[role] diff --git a/subiquity/server/controllers/tests/test_filesystem.py b/subiquity/server/controllers/tests/test_filesystem.py index ee2f4d85..3663f608 100644 --- a/subiquity/server/controllers/tests/test_filesystem.py +++ b/subiquity/server/controllers/tests/test_filesystem.py @@ -1537,11 +1537,11 @@ class TestCoreBootInstallMethods(IsolatedAsyncioTestCase): with mock.patch.object( snapdapi, "post_and_wait", new_callable=mock.AsyncMock ) as mocked: - mocked.return_value = { - "encrypted-devices": { + mocked.return_value = snapdapi.SystemActionResponse( + encrypted_devices={ snapdapi.Role.SYSTEM_DATA: "enc-system-data", }, - } + ) await self.fsc.setup_encryption(context=self.fsc.context) # setup_encryption mutates the filesystem model objects to diff --git a/subiquity/server/snapdapi.py b/subiquity/server/snapdapi.py index 8af125cc..5ae38382 100644 --- a/subiquity/server/snapdapi.py +++ b/subiquity/server/snapdapi.py @@ -24,7 +24,7 @@ import attr from subiquity.common.api.client import make_client from subiquity.common.api.defs import Payload, api, path_parameter -from subiquity.common.serialize import Serializer, named_field +from subiquity.common.serialize import NonExhaustive, Serializer, named_field from subiquity.common.types import Change, TaskStatus log = logging.getLogger("subiquity.server.snapdapi") @@ -90,17 +90,11 @@ class Response: status: str -class Role: +class Role(enum.Enum): NONE = "" MBR = "mbr" SYSTEM_BOOT = "system-boot" - SYSTEM_BOOT_IMAGE = "system-boot-image" - SYSTEM_BOOT_SELECT = "system-boot-select" SYSTEM_DATA = "system-data" - SYSTEM_RECOVERY_SELECT = "system-recovery-select" - SYSTEM_SAVE = "system-save" - SYSTEM_SEED = "system-seed" - SYSTEM_SEED_NULL = "system-seed-null" @attr.s(auto_attribs=True) @@ -134,7 +128,7 @@ class VolumeStructure: offset_write: Optional[RelativeOffset] = named_field("offset-write", None) size: int = 0 type: str = "" - role: str = Role.NONE + role: NonExhaustive[Role] = Role.NONE id: Optional[str] = None filesystem: str = "" content: Optional[List[VolumeContent]] = None @@ -232,6 +226,13 @@ class SystemActionRequest: on_volumes: Dict[str, OnVolume] = named_field("on-volumes") +@attr.s(auto_attribs=True) +class SystemActionResponse: + encrypted_devices: Dict[NonExhaustive[Role], str] = named_field( + "encrypted-devices", default=attr.Factory(dict) + ) + + @api class SnapdAPI: serialize_query_args = False