From 261418ab6d453301f0a7f1e39c8c2bd70c3007fb 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  | 10 ++++++++--
 ovirtlago/testlib.py |  4 ++--
 ovirtlago/virt.py    | 30 +++++++++++++++++++++---------
 3 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/ovirtlago/prefix.py b/ovirtlago/prefix.py
index aee9ca9..11cfbc5 100644
--- a/ovirtlago/prefix.py
+++ b/ovirtlago/prefix.py
@@ -170,8 +170,14 @@ 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 []
+        if isinstance(self.virt_env.engine_vm(), list):
+            engine_dists = list(
+                set([host.distro() for host in self.virt_env.engine_vm()])
+            )
+        else:
+            engine_dists = [self.virt_env.engine_vm().distro()] \
+               if self.virt_env.engine_vm() else []
+
         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..07becce 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')
 
@@ -378,7 +383,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: