From 3e28ccefde8041f7b8d2a9ebbe46686abddd67c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Sun, 14 Oct 2018 03:29:30 +0200
Subject: [PATCH] tests: fix cleanup after backup compatibility tests

Allow removing VMs based on multiple prefixes at once. Removing them
separately doesn't handle all the dependencies (default_netvm, netvm)
correctly. This is needed for backup compatibility tests, where VMs are
created with `test-` prefix and `disp-tests-`. Additionally backup code
will create `disp-no-netvm`, which also may need to be removed.
---
 qubes/tests/__init__.py                  | 18 +++++++++++++-----
 qubes/tests/integ/backupcompatibility.py |  6 ++++--
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/qubes/tests/__init__.py b/qubes/tests/__init__.py
index 32bc888b..6bdf5977 100644
--- a/qubes/tests/__init__.py
+++ b/qubes/tests/__init__.py
@@ -902,9 +902,16 @@ class SystemTestCase(QubesTestCase):
             self._remove_vm_qubes(vm)
 
     def remove_test_vms(self, xmlpath=XMLPATH, prefix=VMPREFIX):
-        '''Aggresively remove any domain that has name in testing namespace.
+        '''Aggressively remove any domain that has name in testing namespace.
+
+        :param prefix: name prefix of VMs to remove, can be a list of prefixes
         '''
 
+        if isinstance(prefix, str):
+            prefixes = [prefix]
+        else:
+            prefixes = prefix
+        del prefix
         # first, remove them Qubes-way
         if os.path.exists(xmlpath):
             try:
@@ -917,7 +924,7 @@ class SystemTestCase(QubesTestCase):
                 except AttributeError:
                     host_app = qubes.Qubes()
                 self.remove_vms([vm for vm in app.domains
-                    if vm.name.startswith(prefix) or
+                    if any(vm.name.startswith(prefix) for prefix in prefixes) or
                        (isinstance(vm, qubes.vm.dispvm.DispVM) and vm.name
                         not in host_app.domains)])
                 if not hasattr(self, 'host_app'):
@@ -933,7 +940,7 @@ class SystemTestCase(QubesTestCase):
         # now remove what was only in libvirt
         conn = libvirt.open(qubes.config.defaults['libvirt_uri'])
         for dom in conn.listAllDomains():
-            if dom.name().startswith(prefix):
+            if any(dom.name().startswith(prefix) for prefix in prefixes):
                 self._remove_vm_libvirt(dom)
         conn.close()
 
@@ -948,11 +955,12 @@ class SystemTestCase(QubesTestCase):
             if not os.path.exists(dirpath):
                 continue
             for name in os.listdir(dirpath):
-                if name.startswith(prefix):
+                if any(name.startswith(prefix) for prefix in prefixes):
                     vmnames.add(name)
         for vmname in vmnames:
             self._remove_vm_disk(vmname)
-        self._remove_vm_disk_lvm(prefix)
+        for prefix in prefixes:
+            self._remove_vm_disk_lvm(prefix)
 
     def qrexec_policy(self, service, source, destination, allow=True,
             action=None):
diff --git a/qubes/tests/integ/backupcompatibility.py b/qubes/tests/integ/backupcompatibility.py
index 53165d2c..c7a7f6c4 100644
--- a/qubes/tests/integ/backupcompatibility.py
+++ b/qubes/tests/integ/backupcompatibility.py
@@ -123,8 +123,10 @@ class TC_00_BackupCompatibility(
         qubes.tests.integ.backup.BackupTestsMixin, qubes.tests.SystemTestCase):
 
     def tearDown(self):
-        self.remove_test_vms(prefix="test-")
-        self.remove_test_vms(prefix="disp-test-")
+        prefixes = ["test-", "disp-test-"]
+        if 'disp-no-netvm' not in self.host_app.domains:
+            prefixes.append('disp-no-netvm')
+        self.remove_test_vms(prefix=prefixes)
         super(TC_00_BackupCompatibility, self).tearDown()
 
     def create_whitelisted_appmenus(self, filename):