Skip to content

Commit

Permalink
q-dev: updated tests and wait for attachment to be done
Browse files Browse the repository at this point in the history
  • Loading branch information
piotrbartman committed Oct 14, 2024
1 parent a80d27d commit 094af98
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 42 deletions.
6 changes: 5 additions & 1 deletion qubesusbproxy/core3ext.py
Original file line number Diff line number Diff line change
Expand Up @@ -711,8 +711,12 @@ async def on_domain_start(self, vm, _event, **_kwargs):
if device not in to_attach:
# make it unique
to_attach[device] = assignment.clone(device=device)
in_progress = set()
for assignment in to_attach.values():
asyncio.ensure_future(self.attach_and_notify(vm, assignment))
in_progress.add(
asyncio.ensure_future(self.attach_and_notify(vm, assignment)))
if in_progress:
await asyncio.wait(in_progress)

@qubes.ext.handler('domain-shutdown')
async def on_domain_shutdown(self, vm, _event, **_kwargs):
Expand Down
88 changes: 48 additions & 40 deletions qubesusbproxy/tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
#
import unittest
from unittest import mock
from unittest.mock import Mock
from unittest.mock import Mock, AsyncMock

import jinja2

Expand All @@ -40,10 +40,10 @@
try:
from qubes.device_protocol import DeviceAssignment, VirtualDevice, Port

def make_assignment(backend, ident, required=False):
def make_assignment(backend, ident, auto_attach=False):
return DeviceAssignment(VirtualDevice(Port(
backend, ident, 'usb')),
mode='required' if required else 'manual')
mode='auto-attach' if auto_attach else 'manual')

def assign(test, collection, assignment):
test.loop.run_until_complete(collection.assign(assignment))
Expand Down Expand Up @@ -401,7 +401,8 @@ def setUp(self):
self.qrexec_policy('qubes.USB', self.frontend.name, self.backend.name)
self.usbdev_ident = create_usb_gadget(self.backend).decode()
self.usbdev_name = '{}:{}:{}'.format(
self.backend.name, self.usbdev_ident, "0000:0000::?******")
self.backend.name, self.usbdev_ident,
"1234:1234:0123456789:u080650")

def tearDown(self):
# remove vms in this specific order, otherwise there may remain stray
Expand All @@ -416,9 +417,7 @@ def test_000_list(self):

def test_010_assign(self):
usb_dev = self.backend.devices['usb'][self.usbdev_ident]
ass = DeviceAssignment(VirtualDevice(Port(
self.backend, self.usbdev_ident, 'usb')),
mode='ask-to-attach')
ass = make_assignment(self.backend, self.usbdev_ident, auto_attach=True)
assign(self, self.frontend.devices['usb'], ass)
self.assertIsNone(usb_dev.attachment)
try:
Expand All @@ -432,10 +431,14 @@ def test_010_assign(self):

self.assertEqual(usb_dev.attachment, self.frontend)

@unittest.mock.patch('qubes.ext.utils.confirm_device_attachment')
@unittest.skipIf(LEGACY, "new feature")
def test_011_assign_ask(self):
def test_011_assign_ask(self, confirm):
confirm.return_value = self.frontend.name
usb_dev = self.backend.devices['usb'][self.usbdev_ident]
ass = make_assignment(self.backend, self.usbdev_ident, required=True)
ass = DeviceAssignment(VirtualDevice(Port(
self.backend, self.usbdev_ident, 'usb')),
mode='ask-to-attach')
assign(self, self.frontend.devices['usb'], ass)
self.assertIsNone(usb_dev.attachment)
try:
Expand Down Expand Up @@ -488,7 +491,7 @@ def test_030_detach(self):

def test_040_unassign(self):
usb_dev = self.backend.devices['usb'][self.usbdev_ident]
ass = make_assignment(self.backend, self.usbdev_ident, required=True)
ass = make_assignment(self.backend, self.usbdev_ident, auto_attach=True)
assign(self, self.frontend.devices['usb'], ass)
self.assertIsNone(usb_dev.attachment)
unassign(self, self.frontend.devices['usb'], ass)
Expand Down Expand Up @@ -540,7 +543,7 @@ def test_060_auto_detach_on_remove(self):
def test_061_auto_attach_on_reconnect(self):
self.frontend.start()
usb_list = self.backend.devices['usb']
ass = make_assignment(self.backend, self.usbdev_ident, required=True)
ass = make_assignment(self.backend, self.usbdev_ident, auto_attach=True)
try:
assign(self, self.frontend.devices['usb'], ass)
except qubesusbproxy.core3ext.USBProxyNotInstalled as e:
Expand Down Expand Up @@ -568,7 +571,7 @@ def test_061_auto_attach_on_reconnect(self):
def test_062_ask_to_attach_on_start(self):
self.frontend.start()
usb_list = self.backend.devices['usb']
ass = make_assignment(self.backend, self.usbdev_ident, required=True)
ass = make_assignment(self.backend, self.usbdev_ident, auto_attach=True)
try:
assign(self, self.frontend.devices['usb'], ass)
except qubesusbproxy.core3ext.USBProxyNotInstalled as e:
Expand Down Expand Up @@ -855,8 +858,10 @@ def test_012_on_qdb_change_multiple_assignments_dev(self):
self.assertEqual(self.ext.attach_and_notify.call_args[0][1].options,
{'any': 'did'})

@unittest.mock.patch('subprocess.Popen')
def test_013_on_qdb_change_two_fronts_failed(self, mock_confirm):
# replace async function with sync one
@unittest.mock.patch('qubes.ext.utils.confirm_device_attachment',
new_callable=Mock)
def test_013_on_qdb_change_two_fronts_failed(self, _confirm):
back, front = self.added_assign_setup()

exp_dev = qubesusbproxy.core3ext.USBDevice(back, '1-1')
Expand All @@ -866,18 +871,22 @@ def test_013_on_qdb_change_two_fronts_failed(self, mock_confirm):
back.devices['usb']._assigned.append(assign)
back.devices['usb']._exposed.append(exp_dev)

proc = Mock()
proc.communicate = Mock()
proc.communicate.return_value = (b'nonsense', b'')
mock_confirm.return_value = proc
class quick_mock:
@staticmethod
def result():
return "nonsense"

self.ext.attach_and_notify = Mock()
with mock.patch('asyncio.ensure_future'):
with mock.patch('asyncio.ensure_future') as future:
future.return_value = quick_mock
self.ext.on_qdb_change(back, None, None)
proc.communicate.assert_called_once()

self.ext.attach_and_notify.assert_not_called()

@unittest.mock.patch('subprocess.Popen')
def test_014_on_qdb_change_two_fronts(self, mock_confirm):
# replace async function with sync one
@unittest.mock.patch('qubes.ext.utils.confirm_device_attachment',
new_callable=Mock)
def test_014_on_qdb_change_two_fronts(self, _confirm):
back, front = self.added_assign_setup()

exp_dev = qubesusbproxy.core3ext.USBDevice(back, '1-1')
Expand All @@ -887,16 +896,17 @@ def test_014_on_qdb_change_two_fronts(self, mock_confirm):
back.devices['usb']._assigned.append(assign)
back.devices['usb']._exposed.append(exp_dev)

proc = Mock()
proc.communicate = Mock()
proc.communicate.return_value = (b'front-vm', b'')
mock_confirm.return_value = proc
class quick_mock:
@staticmethod
def result():
return "front-vm"

self.ext.attach_and_notify = Mock()
with mock.patch('asyncio.ensure_future'):
with mock.patch('asyncio.ensure_future') as future:
future.return_value = quick_mock
self.ext.on_qdb_change(back, None, None)
proc.communicate.assert_called_once()
self.ext.attach_and_notify.assert_called_once_with(
front, assign)

self.ext.attach_and_notify.assert_called_once_with(front, assign)
# don't ask again
self.assertEqual(self.ext.attach_and_notify.call_args[0][1].mode.value,
'auto-attach')
Expand All @@ -915,6 +925,7 @@ def test_015_on_qdb_change_ask(self):
self.ext.on_qdb_change(back, None, None)
self.assertEqual(self.ext.attach_and_notify.call_args[0][1].mode.value,
'ask-to-attach')

def test_020_on_startup_multiple_assignments_including_full(self):
back, front = self.added_assign_setup()

Expand All @@ -936,10 +947,9 @@ def test_020_on_startup_multiple_assignments_including_full(self):
back.devices['usb']._exposed.append(
qubesusbproxy.core3ext.USBDevice(back, '1-1'))

self.ext.attach_and_notify = Mock()
self.ext.attach_and_notify = AsyncMock()
loop = asyncio.get_event_loop()
with mock.patch('asyncio.ensure_future'):
loop.run_until_complete(self.ext.on_domain_start(front, None))
loop.run_until_complete(self.ext.on_domain_start(front, None))
self.assertEqual(self.ext.attach_and_notify.call_args[0][1].options,
{'pid': 'did'})

Expand All @@ -961,9 +971,8 @@ def test_021_on_startup_multiple_assignments_port_vs_dev(self):
qubesusbproxy.core3ext.USBDevice(back, '1-1'))

loop = asyncio.get_event_loop()
self.ext.attach_and_notify = Mock()
with mock.patch('asyncio.ensure_future'):
loop.run_until_complete(self.ext.on_domain_start(front, None))
self.ext.attach_and_notify = AsyncMock()
loop.run_until_complete(self.ext.on_domain_start(front, None))
self.assertEqual(self.ext.attach_and_notify.call_args[0][1].options,
{'pid': 'any'})

Expand All @@ -986,10 +995,9 @@ def test_022_on_startup_multiple_assignments_dev(self):
back.devices['usb']._exposed.append(
qubesusbproxy.core3ext.USBDevice(back, '1-2'))

self.ext.attach_and_notify = Mock()
self.ext.attach_and_notify = AsyncMock()
loop = asyncio.get_event_loop()
with mock.patch('asyncio.ensure_future'):
loop.run_until_complete(self.ext.on_domain_start(front, None))
loop.run_until_complete(self.ext.on_domain_start(front, None))
self.assertEqual(self.ext.attach_and_notify.call_args[0][1].options,
{'any': 'did'})

Expand All @@ -1001,7 +1009,7 @@ def test_023_on_startup_already_attached(self):
exp_dev.port, exp_dev.device_id), mode='auto-attach')

front.devices['usb']._assigned.append(assign)
attached_device = back.devices['usb']._exposed.append(exp_dev)
back.devices['usb']._exposed.append(exp_dev)

self.ext.attach_and_notify = Mock()
loop = asyncio.get_event_loop()
Expand Down
1 change: 0 additions & 1 deletion qubesusbproxy/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301,
# USA.
import asyncio
import subprocess
import sys

import qubes
Expand Down

0 comments on commit 094af98

Please sign in to comment.