zdev: use attrs for storing actions internally

Signed-off-by: Olivier Gayot <[email protected]>

subiquity/server/controllers/tests/test_zdev.py

@@ -16,7 +16,7 @@
 import unittest
 from unittest.mock import AsyncMock, Mock, patch
 
-from subiquity.server.controllers.zdev import ZdevController
+from subiquity.server.controllers.zdev import ZdevAction, ZdevController
 from subiquitycore.tests.mocks import make_app
 
 
@@ -25,18 +25,29 @@ def setUp(self):
         self.ctrler = ZdevController(make_app())
 
     def test_make_autoinstall_no_dupes(self):
-        self.ctrler.done_actions = [
+        self.ctrler.done_ai_actions = [
+            ZdevAction(id="0.0.1507", enable=True),
+            ZdevAction(id="0.0.1508", enable=False),
+            ZdevAction(id="0.0.1509", enable=True),
+        ]
+
+        expected = [
             {"id": "0.0.1507", "enabled": True},
             {"id": "0.0.1508", "enabled": False},
             {"id": "0.0.1509", "enabled": True},
         ]
-        self.assertEqual(self.ctrler.done_actions, self.ctrler.make_autoinstall())
+        self.assertEqual(expected, self.ctrler.make_autoinstall())
 
     def test_make_autoinstall_with_dupes(self):
-        action1 = {"id": "0.0.1507", "enabled": True}
-        action2 = {"id": "0.0.1508", "enabled": True}
-        self.ctrler.done_actions = [action1, action1, action1, action2, action2]
-        self.assertEqual([action1, action2], self.ctrler.make_autoinstall())
+        action1 = ZdevAction(id="0.0.1507", enable=True)
+        action2 = ZdevAction(id="0.0.1508", enable=True)
+        self.ctrler.done_ai_actions = [action1, action1, action1, action2, action2]
+
+        expected = [
+            {"id": "0.0.1507", "enabled": True},
+            {"id": "0.0.1508", "enabled": True},
+        ]
+        self.assertEqual(expected, self.ctrler.make_autoinstall())
 
     @patch("asyncio.sleep", AsyncMock())
     async def test_handle_zdevs__none(self):
@@ -67,14 +78,14 @@ async def test_handle_zdevs(self):
         self.assertEqual(expected_calls, m_chzdev.mock_calls)
 
     async def test_chzdev_wrong_action(self):
-        self.ctrler.done_actions = []
+        self.ctrler.done_ai_actions = []
         with self.assertRaises(ValueError):
             await self.ctrler.chzdev("enAble", self.ctrler.zdevinfos["0.0.1507"])
-        self.assertFalse(self.ctrler.done_actions)
+        self.assertFalse(self.ctrler.done_ai_actions)
 
     @patch("asyncio.sleep", AsyncMock())
     async def test_chzdev_enable(self):
-        self.ctrler.done_actions = []
+        self.ctrler.done_ai_actions = []
 
         self.ctrler.app.command_runner = Mock()
         with patch.object(self.ctrler.app.command_runner, "run", AsyncMock()) as m_run:
@@ -83,26 +94,26 @@ async def test_chzdev_enable(self):
         m_run.assert_called_once_with(["chzdev", "--enable", "0.0.1507"])
 
         self.assertEqual(
-            [{"id": "0.0.1507", "state": "enabled"}], self.ctrler.done_actions
+            [ZdevAction(id="0.0.1507", enable=True)], self.ctrler.done_ai_actions
         )
 
     @patch("asyncio.sleep", AsyncMock())
     async def test_chzdev_disable(self):
-        self.ctrler.done_actions = []
+        self.ctrler.done_ai_actions = []
 
         self.ctrler.app.command_runner = Mock()
         with patch.object(self.ctrler.app.command_runner, "run", AsyncMock()) as m_run:
             await self.ctrler.chzdev("disable", self.ctrler.zdevinfos["0.0.1507"])
 
         self.assertEqual(
-            [{"id": "0.0.1507", "state": "disabled"}], self.ctrler.done_actions
+            [ZdevAction(id="0.0.1507", enable=False)], self.ctrler.done_ai_actions
         )
 
         m_run.assert_called_once_with(["chzdev", "--disable", "0.0.1507"])
 
     @patch("asyncio.sleep", AsyncMock())
     async def test_chzdev_enable_disable_multiple(self):
-        self.ctrler.done_actions = []
+        self.ctrler.done_ai_actions = []
 
         self.ctrler.app.command_runner = Mock()
         with patch.object(self.ctrler.app.command_runner, "run", AsyncMock()) as m_run:
@@ -121,10 +132,10 @@ async def test_chzdev_enable_disable_multiple(self):
 
         self.assertEqual(
             [
-                {"id": "0.0.1507", "state": "enabled"},
-                {"id": "0.0.1507", "state": "enabled"},
-                {"id": "0.0.1508", "state": "disabled"},
-                {"id": "0.0.1508", "state": "enabled"},
+                ZdevAction(id="0.0.1507", enable=True),
+                ZdevAction(id="0.0.1507", enable=True),
+                ZdevAction(id="0.0.1508", enable=False),
+                ZdevAction(id="0.0.1508", enable=True),
             ],
-            self.ctrler.done_actions,
+            self.ctrler.done_ai_actions,
         )

subiquity/server/controllers/zdev.py

@@ -20,6 +20,8 @@
 from collections import OrderedDict
 from typing import List, Literal, Optional, TypedDict
 
+import attrs
+
 from subiquity.common.apidef import API
 from subiquity.common.types import ZdevInfo
 from subiquity.common.types.storage import Bootloader
@@ -603,6 +605,19 @@ class ZdevAiItem(TypedDict, total=True):
 ZdevAi = list[ZdevAiItem]
 
 
+@attrs.define(auto_attribs=True)
+class ZdevAction:
+    id: str
+    enable: bool
+
+    @classmethod
+    def from_ai_item(cls, item: ZdevAiItem) -> "ZdevAction":
+        return cls(id=item["id"], enable=item["enabled"])
+
+    def to_ai_item(self) -> ZdevAiItem:
+        return ZdevAiItem(id=self.id, enabled=self.enable)
+
+
 class ZdevController(SubiquityController):
     endpoint = API.zdev
 
@@ -620,12 +635,12 @@ class ZdevController(SubiquityController):
     autoinstall_default = []
 
     def __init__(self, app):
-        self.ai_data: ZdevAi = []
+        self.ai_actions: list[ZdevAction] = []
         self.zdev_handling_task: Optional[asyncio.Task] = None
 
         # Recording of actions that have been performed on the Zdevs. This is
         # only used to produce an autoinstall config at the end.
-        self.done_actions: list[ZdevAiItem] = []
+        self.done_ai_actions: list[ZdevAction] = []
 
         super().__init__(app)
         if self.opts.dry_run:
@@ -638,26 +653,26 @@ def __init__(self, app):
             self.zdevinfos = OrderedDict([(i.id, i) for i in zdevinfos])
 
     def load_autoinstall_data(self, data: ZdevAi) -> None:
-        self.ai_data = data
+        self.ai_actions = [ZdevAction.from_ai_item(item) for item in data]
 
     @with_context()
     async def apply_autoinstall_config(self, context) -> None:
         if self.zdev_handling_task is not None:
             await self.zdev_handling_task
 
     def start(self) -> None:
-        if self.ai_data:
+        if self.ai_actions:
             self.zdev_handling_task = schedule_task(self.handle_zdevs())
 
     def make_autoinstall(self) -> ZdevAi:
         # Small "optimization" to avoid producing a config that enables or
         # disables a given device multiple times in a row.
-        return [x for x, _ in itertools.groupby(self.done_actions)]
+        return [x.to_ai_item() for x, _ in itertools.groupby(self.done_ai_actions)]
 
     async def handle_zdevs(self) -> None:
-        for item in self.ai_data:
-            action = "enable" if item["enabled"] else "disable"
-            await self.chzdev(action, self.zdevinfos[item["id"]])
+        for ai_action in self.ai_actions:
+            action = "enable" if ai_action.enable else "disable"
+            await self.chzdev(action, self.zdevinfos[ai_action.id])
 
     def interactive(self):
         if self.app.base_model.filesystem.bootloader != Bootloader.NONE:
@@ -669,14 +684,12 @@ async def chzdev(
     ) -> None:
         if action == "enable":
             on = True
-            state = "enabled"
         elif action == "disable":
             on = False
-            state = "disabled"
         else:
             raise ValueError("action must be 'enable' or 'disable'")
 
-        self.done_actions.append(ZdevAiItem(id=zdev.id, state=state))
+        self.done_ai_actions.append(ZdevAction(id=zdev.id, enable=on))
 
         if self.opts.dry_run:
             self.zdevinfos[zdev.id].on = on