From c56ea0ec3bc1e6bc6ee75723e5cdaa194319c901 Mon Sep 17 00:00:00 2001
From: Wojciech Kozlowski <wk@wojciechkozlowski.eu>
Date: Sat, 11 Feb 2023 10:30:17 +0100
Subject: [PATCH] Replace testing shell scripts with python

---
 requirements.txt          |  1 +
 scripts/testing/revert.sh |  5 ---
 scripts/testing/start.sh  |  5 ---
 scripts/testing/stop.sh   |  5 ---
 scripts/testing/vmgr.py   | 95 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 96 insertions(+), 15 deletions(-)
 delete mode 100755 scripts/testing/revert.sh
 delete mode 100755 scripts/testing/start.sh
 delete mode 100755 scripts/testing/stop.sh
 create mode 100755 scripts/testing/vmgr.py

diff --git a/requirements.txt b/requirements.txt
index 069c0f8..091f1ad 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 ansible
 keyring
+libvirt-python
 requests
diff --git a/scripts/testing/revert.sh b/scripts/testing/revert.sh
deleted file mode 100755
index fe7fef2..0000000
--- a/scripts/testing/revert.sh
+++ /dev/null
@@ -1,5 +0,0 @@
-#!/usr/bin/env bash
-
-virsh --connect=qemu:///system snapshot-revert heimdall-virt 00-hostname
-virsh --connect=qemu:///system snapshot-revert valkyrie-virt 00-hostname
-virsh --connect=qemu:///system snapshot-revert yggdrasil-virt 01-hpool
diff --git a/scripts/testing/start.sh b/scripts/testing/start.sh
deleted file mode 100755
index ccc1312..0000000
--- a/scripts/testing/start.sh
+++ /dev/null
@@ -1,5 +0,0 @@
-#!/usr/bin/env bash
-
-virsh --connect=qemu:///system start heimdall-virt
-virsh --connect=qemu:///system start valkyrie-virt
-virsh --connect=qemu:///system start yggdrasil-virt
diff --git a/scripts/testing/stop.sh b/scripts/testing/stop.sh
deleted file mode 100755
index c46104e..0000000
--- a/scripts/testing/stop.sh
+++ /dev/null
@@ -1,5 +0,0 @@
-#!/usr/bin/env bash
-
-virsh --connect=qemu:///system shutdown heimdall-virt
-virsh --connect=qemu:///system shutdown valkyrie-virt
-virsh --connect=qemu:///system shutdown yggdrasil-virt
diff --git a/scripts/testing/vmgr.py b/scripts/testing/vmgr.py
new file mode 100755
index 0000000..6cc8c86
--- /dev/null
+++ b/scripts/testing/vmgr.py
@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+
+import argparse
+import contextlib
+import libvirt
+from typing import Dict, Iterable
+
+CONNECTION = "qemu:///system"
+INVENTORY = ["heimdall-virt", "valkyrie-virt", "yggdrasil-virt"]
+
+
+class VirtDomain:
+
+    __ANSIBLE_EDDA_SNAPSHOT = "ansible-edda"
+
+    def __init__(self, libvirt_domain: libvirt.virDomain):
+        self.__domain = libvirt_domain
+
+    def start(self) -> None:
+        if not self.__domain.isActive():
+            self.__domain.create()
+
+    def stop(self) -> None:
+        if self.__domain.isActive():
+            self.__domain.shutdown()
+
+    def revert(self) -> None:
+        if self.__domain.isActive():
+            raise RuntimeError(f"\"{self.__domain.name()}\" must be stopped before reverting")
+        snap = self.__domain.snapshotLookupByName(self.__ANSIBLE_EDDA_SNAPSHOT)
+        self.__domain.revertToSnapshot(snap)
+
+
+class VirtManager:
+
+    def __init__(self, conn_name: str, hosts: Iterable[str]):
+        self.__conn = libvirt.open(conn_name)
+        self.__inventory = {}
+        for hostname in hosts:
+            try:
+                self.__inventory[hostname] = VirtDomain(self.__conn.lookupByName(hostname))
+            except libvirt.libvirtError as e:
+                if e.get_error_code() != libvirt.VIR_ERR_NO_DOMAIN:
+                    raise
+
+    def close(self):
+        self.__conn.close()
+
+    @property
+    def inventory(self) -> Dict[str, VirtDomain]:
+        return self.__inventory
+
+    def start(self):
+        for vdom in self.__inventory.values():
+            vdom.start()
+
+    def stop(self):
+        for vdom in self.__inventory.values():
+            vdom.stop()
+
+    def revert(self):
+        for vdom in self.__inventory.values():
+            vdom.revert()
+
+
+@contextlib.contextmanager
+def virt_manager(connection_name: str, hosts: Iterable[str]):
+    vmgr = VirtManager(connection_name, hosts)
+    yield vmgr
+    vmgr.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Manage virtual machines for testing")
+
+    parser.add_argument("--connect", type=str, default=CONNECTION,
+                        help="hypervisor connection URI")
+    parser.add_argument("--limit", type=str, default=','.join(INVENTORY),
+                        help="limit to selected hosts")
+
+    subparsers = parser.add_subparsers()
+
+    start_parser = subparsers.add_parser("start")
+    start_parser.set_defaults(func=VirtManager.start)
+
+    stop_parser = subparsers.add_parser("stop")
+    stop_parser.set_defaults(func=VirtManager.stop)
+
+    revert_parser = subparsers.add_parser("revert")
+    revert_parser.set_defaults(func=VirtManager.revert)
+
+    args = parser.parse_args()
+
+    with virt_manager(CONNECTION, args.limit.split(',')) as vmgr:
+        args.func(vmgr)