diff --git a/qubesadmin/devices.py b/qubesadmin/devices.py index f0d33325..673f1f75 100644 --- 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 index 53b108f7..e1f898d2 100644 --- 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,10 +289,10 @@ 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( @@ -299,13 +300,13 @@ def test_070_update_assignment(self): 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( @@ -313,13 +314,13 @@ def test_071_update_assignment_false(self): 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):