appmenus: avoid duplicate dispvm menu updates

Author: Jayant-kernel

qubesappmenus/tests.py

@@ -26,6 +26,7 @@
 import os
 import tempfile
 import sys
+import types
 
 import unittest
 import unittest.mock
@@ -35,6 +36,28 @@
 import qubesappmenus
 import qubesappmenus.receive
 
+try:
+    import qubesappmenusext
+except ModuleNotFoundError as e:
+    if e.name.split('.')[0] != 'qubes':
+        raise
+    qubes_module = types.ModuleType('qubes')
+    qubes_ext_module = types.ModuleType('qubes.ext')
+    qubes_utils_module = types.ModuleType('qubes.utils')
+
+    def handler(*args, **kwargs):
+        # Only decorator registration is needed for these unit tests.
+        return lambda func: func
+
+    qubes_ext_module.handler = handler
+    qubes_ext_module.Extension = object
+    qubes_utils_module.sanitize_stderr_for_log = lambda stderr: stderr
+    qubes_module.ext = qubes_ext_module
+    sys.modules['qubes'] = qubes_module
+    sys.modules['qubes.ext'] = qubes_ext_module
+    sys.modules['qubes.utils'] = qubes_utils_module
+    import qubesappmenusext
+
 class Label(object):
     def __init__(self, index, color, name):
         self.index = index
@@ -48,6 +71,12 @@ class TestApp(object):
     def __init__(self):
         self.domains = {}
 
+class TestVMM(object):
+    offline_mode = False
+
+class TestAppmenusExtApp(object):
+    vmm = TestVMM()
+
 class TestFeatures(dict):
 
     def __init__(self, vm, **kwargs) -> None:
@@ -95,6 +124,18 @@ def icon(self):
             return 'servicevm-' + raw_icon_name
         return 'appvm-' + raw_icon_name
 
+class TestAppmenusExtVM(object):
+    # pylint: disable=too-few-public-methods
+    app = TestAppmenusExtApp()
+
+    def __init__(self, template_for_dispvms=False,
+            appmenus_dispvm=False):
+        self.name = 'test-ext-vm'
+        self.template_for_dispvms = template_for_dispvms
+        self.features = TestFeatures(self)
+        if appmenus_dispvm:
+            self.features['appmenus-dispvm'] = '1'
+
 VMPREFIX = 'test-'
 
 class TC_00_Appmenus(unittest.TestCase):
@@ -146,6 +187,59 @@ def assertPathNotExists(self, path):
         if os.path.exists(path):
             self.fail("Path {} exists while it should not".format(path))
 
+    def assertUpdateScheduled(self, callback, should_schedule):
+        ext = qubesappmenusext.AppmenusExtension()
+        ext.collect_done_tasks = unittest.mock.Mock()
+        ext.update_appmenus = unittest.mock.Mock(return_value='update-task')
+        with unittest.mock.patch('asyncio.ensure_future') as ensure_future:
+            callback(ext)
+        if should_schedule:
+            ext.collect_done_tasks.assert_called_once()
+            ext.update_appmenus.assert_called_once()
+            ensure_future.assert_called_once_with('update-task')
+        else:
+            ext.collect_done_tasks.assert_not_called()
+            ext.update_appmenus.assert_not_called()
+            ensure_future.assert_not_called()
+
+    def test_000_appmenus_ext_template_for_dispvms_needs_feature(self):
+        vm = TestAppmenusExtVM(
+            template_for_dispvms=True,
+            appmenus_dispvm=False)
+
+        self.assertUpdateScheduled(
+            lambda ext: ext.template_for_dispvms_setter(vm, None),
+            False)
+
+        vm.features['appmenus-dispvm'] = '1'
+        self.assertUpdateScheduled(
+            lambda ext: ext.template_for_dispvms_setter(vm, None),
+            True)
+
+    def test_000_appmenus_ext_dispvm_feature_needs_property(self):
+        vm = TestAppmenusExtVM(
+            template_for_dispvms=False,
+            appmenus_dispvm=True)
+
+        self.assertUpdateScheduled(
+            lambda ext: ext.on_feature_set_appmenus_dispvm(
+                vm, None, 'appmenus-dispvm', '1'),
+            False)
+        self.assertUpdateScheduled(
+            lambda ext: ext.on_feature_del_appmenus_dispvm(
+                vm, None, 'appmenus-dispvm'),
+            False)
+
+        vm.template_for_dispvms = True
+        self.assertUpdateScheduled(
+            lambda ext: ext.on_feature_set_appmenus_dispvm(
+                vm, None, 'appmenus-dispvm', '1'),
+            True)
+        self.assertUpdateScheduled(
+            lambda ext: ext.on_feature_del_appmenus_dispvm(
+                vm, None, 'appmenus-dispvm'),
+            True)
+
 
     def test_000_templates_dirs(self):
         self.assertEqual(

qubesappmenusext/__init__.py

@@ -175,6 +175,8 @@ def provides_network_setter(self, vm, event, **kwargs):
     def template_for_dispvms_setter(self, vm, event, **kwargs):
         if vm.app.vmm.offline_mode:
             return
+        if not vm.features.get('appmenus-dispvm', False):
+            return
         self.collect_done_tasks(vm)
         self.vm_tasks[vm.name].append(
             asyncio.ensure_future(self.update_appmenus(vm)))
@@ -194,6 +196,8 @@ def provides_network_setter(self, vm, event, name, newvalue, oldvalue=None):
     def on_feature_del_appmenus_dispvm(self, vm, event, feature):
         if vm.app.vmm.offline_mode:
             return
+        if not getattr(vm, 'template_for_dispvms', False):
+            return
         self.collect_done_tasks(vm)
         self.vm_tasks[vm.name].append(
             asyncio.ensure_future(self.update_appmenus(vm)))
@@ -203,6 +207,8 @@ def on_feature_set_appmenus_dispvm(self, vm, event, feature,
             value, oldvalue=None):
         if vm.app.vmm.offline_mode:
             return
+        if not getattr(vm, 'template_for_dispvms', False):
+            return
         self.collect_done_tasks(vm)
         self.vm_tasks[vm.name].append(
             asyncio.ensure_future(self.update_appmenus(vm)))