use NonExhaustive to make snapdapi.Role an enum again

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]

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

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