From 02c601b8307c1704a4a9e9b5332a73e1b0e13a12 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Thu, 3 Mar 2016 01:09:16 +0100
Subject: [PATCH] tests: expose API for external tests

Fixes QubesOS/qubes-issues#1800
---
 tests/__init__.py | 83 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 83 insertions(+)

diff --git a/tests/__init__.py b/tests/__init__.py
index c8a31a6b..02d660e3 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -36,6 +36,7 @@ import unittest.case
 
 import lxml.etree
 import sys
+import pkg_resources
 
 import qubes.backup
 import qubes.qubes
@@ -490,6 +491,12 @@ class SystemTestsMixin(object):
         subprocess.check_call(command)
 
     def shutdown_and_wait(self, vm, timeout=60):
+        """
+
+        :param vm: VM object
+        :param timeout: timeout after which fail the test
+        :return:
+        """
         vm.shutdown()
         while timeout > 0:
             if not vm.is_running():
@@ -733,6 +740,48 @@ class BackupTestsMixin(SystemTestsMixin):
         f.truncate(size)
         f.close()
 
+class ExtraTestMixin(SystemTestsMixin):
+
+    template = None
+
+    def setUp(self):
+        super(ExtraTestMixin, self).setUp()
+        self.qc.unlock_db()
+
+    def create_vms(self, names):
+        """
+        Create AppVMs for the duration of the test. Will be automatically
+        removed after completing the test.
+        :param names: list of VM names to create (each of them will be
+        prefixed with some test specific string)
+        :return: list of created VM objects
+        """
+        self.qc.lock_db_for_writing()
+        self.qc.load()
+        if self.template:
+            template = self.qc.get_vm_by_name(self.template)
+        else:
+            template = self.qc.get_default_template()
+        for vmname in names:
+            vm = self.qc.add_new_vm("QubesAppVm",
+                                    name=self.make_vm_name(vmname),
+                                    template=template)
+            vm.create_on_disk(verbose=False)
+        self.save_and_reload_db()
+
+        # get objects after reload
+        vms = []
+        for vmname in names:
+            vms.append(self.qc.get_vm_by_name(self.make_vm_name(vmname)))
+        return vms
+
+    def enable_network(self):
+        """
+        Enable access to the network. Must be called before creating VMs.
+        """
+        # nothing to do in core2
+        pass
+
 
 def load_tests(loader, tests, pattern):
     # discard any tests from this module, because it hosts base classes
@@ -752,6 +801,40 @@ def load_tests(loader, tests, pattern):
             ):
         tests.addTests(loader.loadTestsFromName(modname))
 
+    for entry in pkg_resources.iter_entry_points('qubes.tests.extra'):
+        for test_case in entry():
+            tests.addTests(loader.loadTestsFromTestCase(
+                type(
+                    entry.name + '_' + test_case.__name__,
+                    (test_case, ExtraTestMixin, QubesTestCase),
+                    {}
+                )
+            ))
+
+
+    try:
+        qc = qubes.qubes.QubesVmCollection()
+        qc.lock_db_for_reading()
+        qc.load()
+        qc.unlock_db()
+        templates = [vm.name for vm in qc.values() if
+                     isinstance(vm, qubes.qubes.QubesTemplateVm)]
+    except OSError:
+        templates = []
+
+    for entry in pkg_resources.iter_entry_points(
+            'qubes.tests.extra.for_template'):
+        for test_case in entry():
+            for template in templates:
+                tests.addTests(loader.loadTestsFromTestCase(
+                    type(
+                        entry.name + '_' + test_case.__name__,
+                        (test_case, ExtraTestMixin,
+                         QubesTestCase),
+                        {'template': template}
+                    )
+                ))
+
     return tests