From ca2ad86438293af6715a9890b168f159ff4d9b9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Bosdonnat?= <cbosdonnat@suse.com>
Date: Thu, 18 Oct 2018 13:32:59 +0200
Subject: [PATCH] Add virt.all_capabilities

In order to get all possible capabilities from a host, the user has to
call virt.capabilities, and then loop over the guests and domains
before calling virt.domain_capabilities for each of them.

This commit embeds all this logic to get them all in a single
virt.all_capabilities call.
---
 salt/modules/virt.py            | 73 +++++++++++++++++++++++++++++++--
 tests/unit/modules/test_virt.py |  2 +-
 2 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/salt/modules/virt.py b/salt/modules/virt.py
index 313181c49e..362c2a68b5 100644
--- a/salt/modules/virt.py
+++ b/salt/modules/virt.py
@@ -5568,11 +5568,76 @@ def _parse_domain_caps(caps):
     return result
 
 
+def _parse_domain_caps(caps):
+    """
+    Parse the XML document of domain capabilities into a structure.
+    """
+    result = {
+        "emulator": caps.find("path").text if caps.find("path") is not None else None,
+        "domain": caps.find("domain").text if caps.find("domain") is not None else None,
+        "machine": caps.find("machine").text
+        if caps.find("machine") is not None
+        else None,
+        "arch": caps.find("arch").text if caps.find("arch") is not None else None,
+    }
+
+
+def all_capabilities(**kwargs):
+    """
+    Return the host and domain capabilities in a single call.
+
+    .. versionadded:: 3001
+
+    :param connection: libvirt connection URI, overriding defaults
+    :param username: username to connect with, overriding defaults
+    :param password: password to connect with, overriding defaults
+
+    CLI Example:
+
+    .. code-block:: bash
+
+        salt '*' virt.all_capabilities
+
+    """
+    conn = __get_conn(**kwargs)
+    try:
+        host_caps = ElementTree.fromstring(conn.getCapabilities())
+        domains = [
+            [
+                (guest.get("arch", {}).get("name", None), key)
+                for key in guest.get("arch", {}).get("domains", {}).keys()
+            ]
+            for guest in [
+                _parse_caps_guest(guest) for guest in host_caps.findall("guest")
+            ]
+        ]
+        flattened = [pair for item in (x for x in domains) for pair in item]
+        result = {
+            "host": {
+                "host": _parse_caps_host(host_caps.find("host")),
+                "guests": [
+                    _parse_caps_guest(guest) for guest in host_caps.findall("guest")
+                ],
+            },
+            "domains": [
+                _parse_domain_caps(
+                    ElementTree.fromstring(
+                        conn.getDomainCapabilities(None, arch, None, domain)
+                    )
+                )
+                for (arch, domain) in flattened
+            ],
+        }
+        return result
+    finally:
+        conn.close()
+
+
 def domain_capabilities(emulator=None, arch=None, machine=None, domain=None, **kwargs):
     """
     Return the domain capabilities given an emulator, architecture, machine or virtualization type.
 
-    .. versionadded:: 2019.2.0
+    .. versionadded:: Fluorine
 
     :param emulator: return the capabilities for the given emulator binary
     :param arch: return the capabilities for the given CPU architecture
@@ -5611,7 +5676,7 @@ def all_capabilities(**kwargs):
     """
     Return the host and domain capabilities in a single call.
 
-    .. versionadded:: 3001
+    .. versionadded:: Neon
 
     :param connection: libvirt connection URI, overriding defaults
     :param username: username to connect with, overriding defaults
@@ -5625,6 +5690,7 @@ def all_capabilities(**kwargs):
 
     """
     conn = __get_conn(**kwargs)
+    result = {}
     try:
         host_caps = ElementTree.fromstring(conn.getCapabilities())
         domains = [
@@ -5653,10 +5719,11 @@ def all_capabilities(**kwargs):
                 for (arch, domain) in flattened
             ],
         }
-        return result
     finally:
         conn.close()
 
+    return result
+
 
 def cpu_baseline(full=False, migratable=False, out="libvirt", **kwargs):
     """
diff --git a/tests/unit/modules/test_virt.py b/tests/unit/modules/test_virt.py
index cce107c9e4..e9e73d7b5d 100644
--- a/tests/unit/modules/test_virt.py
+++ b/tests/unit/modules/test_virt.py
@@ -4063,7 +4063,7 @@ class VirtTestCase(TestCase, LoaderModuleMockMixin):
             "44454c4c-3400-105a-8033-b3c04f4b344a", caps["host"]["host"]["uuid"]
         )
         self.assertEqual(
-            {"qemu", "kvm"}, {domainCaps["domain"] for domainCaps in caps["domains"]},
+            {"qemu", "kvm"}, {domainCaps["domain"] for domainCaps in caps["domains"]}
         )
 
     def test_network_tag(self):
-- 
2.29.2