From 386089c9e5d7b1e18dafddaf1178ff12626c29cd Mon Sep 17 00:00:00 2001
From: Galit <galfrid@gmail.com>
Date: Mon, 3 Sep 2018 10:09:17 +0300
Subject: [PATCH] Add support for multiple engine

Signed-off-by: Galit <galfrid@gmail.com>
---
 ovirtlago/prefix.py  |  6 ++++--
 ovirtlago/testlib.py |  4 ++--
 ovirtlago/virt.py    | 39 ++++++++++++++++++++++++++-------------
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/ovirtlago/prefix.py b/ovirtlago/prefix.py
index aee9ca9..b93499b 100644
--- a/ovirtlago/prefix.py
+++ b/ovirtlago/prefix.py
@@ -170,8 +170,10 @@ def prepare_repo(
         sub_repo_name='default'
     ):
         # Detect distros from template metadata
-        engine_dists = [self.virt_env.engine_vm().distro()] \
-            if self.virt_env.engine_vm() else []
+        engine_dists = list(
+            set([host.distro() for host in self.virt_env.engine_vm()])
+        )
+
         vdsm_dists = list(
             set([host.distro() for host in self.virt_env.host_vms()])
         )
diff --git a/ovirtlago/testlib.py b/ovirtlago/testlib.py
index 68d7150..6996c4b 100644
--- a/ovirtlago/testlib.py
+++ b/ovirtlago/testlib.py
@@ -107,12 +107,12 @@ def get_vm_caps(vm):
     return caps.issubset(existing_caps)
 
 
-def engine_capability(caps):
+def engine_capability(caps, idx=0):
     def decorator(func):
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
             prefix = get_test_prefix()
-            if not _vms_capable([prefix.virt_env.engine_vm()], caps):
+            if not _vms_capable([prefix.virt_env.engine_vm()[idx]], caps):
                 raise SkipTest()
             return func()
 
diff --git a/ovirtlago/virt.py b/ovirtlago/virt.py
index 566c885..f8670bc 100644
--- a/ovirtlago/virt.py
+++ b/ovirtlago/virt.py
@@ -49,7 +49,7 @@
 
 class OvirtVirtEnv(lago.virt.VirtEnv):
     def __init__(self, prefix, vm_specs, net_spec):
-        self._engine_vm = None
+        self._engine_vm = []
         self._host_vms = []
         super(OvirtVirtEnv, self).__init__(prefix, vm_specs, net_spec)
 
@@ -71,12 +71,11 @@ def _create_vm(self, vm_spec):
             )
 
         if provider_name == 'ovirt-engine':
-            if self._engine_vm is not None:
-                raise RuntimeError('Engine VM already exists')
-
             vm_spec['vm-type'] = provider_name
-            self._engine_vm = super(OvirtVirtEnv, self)._create_vm(vm_spec)
-            return self._engine_vm
+            self._engine_vm.append(
+                super(OvirtVirtEnv, self)._create_vm(vm_spec)
+            )
+            return self._engine_vm[-1]
 
         elif provider_name in ('ovirt-host', 'ovirt-node'):
             vm_spec['vm-type'] = provider_name
@@ -89,7 +88,13 @@ def _create_vm(self, vm_spec):
             return super(OvirtVirtEnv, self)._create_vm(vm_spec)
 
     def engine_vm(self):
-        return self._engine_vm
+        if len(self._engine_vm) == 1:
+            return self._engine_vm[0]
+        else:
+            return self._engine_vm[:]
+
+    def engine_vms(self):
+        return self._engine_vm[:]
 
     def host_vms(self):
         return self._host_vms[:]
@@ -121,7 +126,7 @@ def get_ovirt_cpu_family(self, host=None):
             except IndexError:
                 pass
         if not host:
-            host = self.engine_vm()
+            host = self.engine_vms()[-1]
         if host is None:
             raise RuntimeError('No Engine or Host VMs found')
 
@@ -154,9 +159,9 @@ def get_ovirt_cpu_family(self, host=None):
         return cpu_map[host.cpu_vendor][host.cpu_model]
 
     @require_sdk(version='4')
-    def update_clusters_cpu(self, timeout=2 * 60):
+    def update_clusters_cpu(self, timeout=2 * 60, idx=0):
         cpu_family = self.get_ovirt_cpu_family()
-        api = self.engine_vm().get_api_v4(check=True)
+        api = self.engine_vm()[idx].get_api_v4(check=True)
         clusters_service = api.system_service().clusters_service()
         clusters = clusters_service.list()
 
@@ -229,7 +234,7 @@ def _vdsm_up(host):
                 ),
             )
 
-    def assert_engine_alive(self, timeout=2 * 60):
+    def assert_engine_alive(self, timeout=2 * 60, idx=0):
         """
         Assert service 'ovirt-engine' reports running on the engine VM
 
@@ -250,7 +255,8 @@ def _ovirt_engine_up(host):
             return status
 
         testlib.assert_true_within(
-            partial(_ovirt_engine_up, self.engine_vm()),
+            partial(_ovirt_engine_up,
+                    self.engine_vm()[idx]),
             timeout=timeout,
             allowed_exceptions=self._get_check_running_allowed_exceptions(),
         )
@@ -378,7 +384,14 @@ def add_iso(self, path):
         if ret:
             raise RuntimeError('Failed to remove uploaded image')
 
-    def engine_setup(self, config=None):
+    @require_sdk(version='4')
+    def engine_setup(self, timeout=8 * 60):
+        api = self.get_api_v4(check=True)
+        vms_service = api.system_service().vms_service()
+        ids = self._search_vms(vms_service, query='name=HostedEngine')
+        [vms_service.vm_service(id).engine_setup_one() for id in ids]
+
+    def engine_setup_one(self, config=None):
         self.wait_for_ssh()
 
         if config: