q-dev: Set.assignment

QubesOS · Oct 14, 2024 · 091596b · 091596b
1 parent 2c58376
commit 091596b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git a/qubesadmin/devices.py b/qubesadmin/devices.py
@@ -36,7 +36,8 @@ class is implemented by an extension.
 
 import qubesadmin.exc
 from qubesadmin.device_protocol import (Port, DeviceInfo, UnknownDevice,
-                                        DeviceAssignment, VirtualDevice)
+                                        DeviceAssignment, VirtualDevice,
+                                        AssignmentMode)
 
 
 DEVICE_DENY_LIST = "/etc/qubes/device-deny.list"
@@ -195,7 +196,7 @@ def get_exposed_devices(self) -> Iterable[DeviceInfo]:
                 expected_devclass=self._class,
             )
 
-    def update_assignment(self, device: Port, required: Optional[bool]):
+    def update_assignment(self, device: Port, required: AssignmentMode):
         """
         Update assignment of already attached device.
 
@@ -208,7 +209,7 @@ def update_assignment(self, device: Port, required: Optional[bool]):
         self._vm.qubesd_call(
             None,
             'admin.vm.device.{}.Set.assignment'.format(self._class),
-            repr(device), repr(required).encode('utf-8')
+            repr(device), required.value.encode('utf-8')
         )
 
     __iter__ = get_exposed_devices

diff --git a/qubesadmin/tests/devices.py b/qubesadmin/tests/devices.py
@@ -21,8 +21,9 @@
 import qubesadmin.tests
 import qubesadmin.device_protocol
 
-from qubesadmin.device_protocol import (DeviceAssignment, Port, VirtualDevice,
-                                        DeviceInfo, UnknownDevice)
+from qubesadmin.device_protocol import (
+    DeviceAssignment, Port, VirtualDevice, DeviceInfo, UnknownDevice,
+    AssignmentMode)
 
 
 serialized_test_device = (
@@ -288,38 +289,38 @@ def test_060_attached(self):
         self.assertEqual(devs[1].port_id, 'dev2')
         self.assertAllCalled()
 
-    def test_070_update_assignment(self):
+    def test_070_update_assignment_required(self):
         self.app.expected_calls[
             ('test-vm', 'admin.vm.device.test.Set.assignment',
-             'test-vm2+dev1:*', b'True')] = b'0\0'
+             'test-vm2+dev1:*', b'required')] = b'0\0'
         dev = DeviceAssignment(
             VirtualDevice(
                 Port(
                     self.app.domains['test-vm2'],
                     devclass='test',
                     port_id='dev1'),
             ))
-        self.vm.devices['test'].update_assignment(dev, True)
+        self.vm.devices['test'].update_assignment(dev, AssignmentMode.REQUIRED)
         self.assertAllCalled()
 
-    def test_071_update_assignment_false(self):
+    def test_071_update_assignment_ask(self):
         self.app.expected_calls[
             ('test-vm', 'admin.vm.device.test.Set.assignment',
-             'test-vm2+dev1:*', b'False')] = b'0\0'
+             'test-vm2+dev1:*', b'ask-to-attach')] = b'0\0'
         dev = DeviceAssignment(
             VirtualDevice(
                 Port(
                     self.app.domains['test-vm2'],
                     devclass='test',
                     port_id='dev1'),
             ))
-        self.vm.devices['test'].update_assignment(dev, False)
+        self.vm.devices['test'].update_assignment(dev, AssignmentMode.ASK)
         self.assertAllCalled()
 
-    def test_072_update_assignment_none(self):
+    def test_072_update_assignment_auto(self):
         self.app.expected_calls[
             ('test-vm', 'admin.vm.device.test.Set.assignment',
-             'test-vm2+dev1:*', b'None')] = b'0\0'
+             'test-vm2+dev1:*', b'auto-attach')] = b'0\0'
         dev = DeviceAssignment(
             VirtualDevice(
                 Port(
@@ -328,7 +329,7 @@ def test_072_update_assignment_none(self):
                     port_id='dev1',
                 )
             ))
-        self.vm.devices['test'].update_assignment(dev, None)
+        self.vm.devices['test'].update_assignment(dev, AssignmentMode.AUTO)
         self.assertAllCalled()
 
     def test_073_list(self):