From 12f18277a8d0752379c727e11fa245a4de0f5781 Mon Sep 17 00:00:00 2001
From: George Matthews <george.matthews@stfc.ac.uk>
Date: Tue, 10 Dec 2024 13:55:07 +0000
Subject: [PATCH] Update hypervisor properties

Catch missing key error instead of defaulting to None
Remove deprecated current workload property
---
 .../enums/props/hypervisor_properties.py      | 56 +++++++++----------
 openstackquery/mappings/hypervisor_mapping.py |  2 +-
 openstackquery/query_blocks/result.py         |  2 +-
 3 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/openstackquery/enums/props/hypervisor_properties.py b/openstackquery/enums/props/hypervisor_properties.py
index b9fa5d3..1517a8c 100644
--- a/openstackquery/enums/props/hypervisor_properties.py
+++ b/openstackquery/enums/props/hypervisor_properties.py
@@ -12,7 +12,7 @@ class HypervisorProperties(PropEnum):
     An enum class for all hypervisor properties
     """
 
-    HYPERVISOR_CURRENT_WORKLOAD = auto()
+    # HYPERVISOR_CURRENT_WORKLOAD = auto()
     HYPERVISOR_DISK_FREE = auto()
     HYPERVISOR_DISK_SIZE = auto()
     HYPERVISOR_DISK_USED = auto()
@@ -35,10 +35,10 @@ def _get_aliases() -> Dict:
         A method that returns all valid string alias mappings
         """
         return {
-            HypervisorProperties.HYPERVISOR_CURRENT_WORKLOAD: [
-                "current_workload",
-                "workload",
-            ],
+            # HypervisorProperties.HYPERVISOR_CURRENT_WORKLOAD: [
+            #     "current_workload",
+            #     "workload",
+            # ],
             HypervisorProperties.HYPERVISOR_DISK_FREE: [
                 "local_disk_free",
                 "free_disk_gb",
@@ -74,39 +74,39 @@ def get_prop_mapping(prop) -> Optional[PropFunc]:
         :param prop: A HypervisorProperty Enum for which a function may exist for
         """
         mapping = {
-            HypervisorProperties.HYPERVISOR_CURRENT_WORKLOAD: lambda a: a[
-                "current_workload"
+            # HypervisorProperties.HYPERVISOR_CURRENT_WORKLOAD: lambda a: a[
+            #     "current_workload"
+            # ],
+            HypervisorProperties.HYPERVISOR_DISK_FREE: lambda a: a.resources["DISK_GB"][
+                "free"
+            ],
+            HypervisorProperties.HYPERVISOR_DISK_SIZE: lambda a: a.resources["DISK_GB"][
+                "total"
+            ],
+            HypervisorProperties.HYPERVISOR_DISK_USED: lambda a: a.resources["DISK_GB"][
+                "usage"
             ],
-            HypervisorProperties.HYPERVISOR_DISK_FREE: lambda a: a.resources.get(
-                "DISK_GB"
-            ).get("free", None),
-            HypervisorProperties.HYPERVISOR_DISK_SIZE: lambda a: a.resources.get(
-                "DISK_GB"
-            ).get("total", None),
-            HypervisorProperties.HYPERVISOR_DISK_USED: lambda a: a.resources.get(
-                "DISK_GB"
-            ).get("usage", None),
             HypervisorProperties.HYPERVISOR_ID: lambda a: a["id"],
             HypervisorProperties.HYPERVISOR_IP: lambda a: a["host_ip"],
-            HypervisorProperties.HYPERVISOR_MEMORY_FREE: lambda a: a.resources.get(
+            HypervisorProperties.HYPERVISOR_MEMORY_FREE: lambda a: a.resources[
                 "MEMORY_MB"
-            ).get("free", None),
-            HypervisorProperties.HYPERVISOR_MEMORY_SIZE: lambda a: a.resources.get(
+            ]["free"],
+            HypervisorProperties.HYPERVISOR_MEMORY_SIZE: lambda a: a.resources[
                 "MEMORY_MB"
-            ).get("total", None),
-            HypervisorProperties.HYPERVISOR_MEMORY_USED: lambda a: a.resources.get(
+            ]["total"],
+            HypervisorProperties.HYPERVISOR_MEMORY_USED: lambda a: a.resources[
                 "MEMORY_MB"
-            ).get("usage", None),
+            ]["usage"],
             HypervisorProperties.HYPERVISOR_NAME: lambda a: a["name"],
             # HypervisorProperties.HYPERVISOR_SERVER_COUNT: lambda a: a["runnning_vms"],
             HypervisorProperties.HYPERVISOR_STATE: lambda a: a["state"],
             HypervisorProperties.HYPERVISOR_STATUS: lambda a: a["status"],
-            HypervisorProperties.HYPERVISOR_VCPUS: lambda a: a.resources.get(
-                "VCPU"
-            ).get("total", None),
-            HypervisorProperties.HYPERVISOR_VCPUS_USED: lambda a: a.resources.get(
-                "VCPU"
-            ).get("usage", None),
+            HypervisorProperties.HYPERVISOR_VCPUS: lambda a: a.resources["VCPU"][
+                "total"
+            ],
+            HypervisorProperties.HYPERVISOR_VCPUS_USED: lambda a: a.resources["VCPU"][
+                "usage"
+            ],
             HypervisorProperties.HYPERVISOR_DISABLED_REASON: lambda a: a["service"][
                 "disabled_reason"
             ],
diff --git a/openstackquery/mappings/hypervisor_mapping.py b/openstackquery/mappings/hypervisor_mapping.py
index 3e01f8e..4ff4344 100644
--- a/openstackquery/mappings/hypervisor_mapping.py
+++ b/openstackquery/mappings/hypervisor_mapping.py
@@ -82,7 +82,7 @@ def get_client_side_handlers() -> QueryClientSideHandlers:
             HypervisorProperties.HYPERVISOR_VCPUS,
             HypervisorProperties.HYPERVISOR_VCPUS_USED,
             # HypervisorProperties.HYPERVISOR_SERVER_COUNT, # Deprecated, use server query
-            HypervisorProperties.HYPERVISOR_CURRENT_WORKLOAD,
+            # HypervisorProperties.HYPERVISOR_CURRENT_WORKLOAD,
         ]
 
         return QueryClientSideHandlers(
diff --git a/openstackquery/query_blocks/result.py b/openstackquery/query_blocks/result.py
index 82eeffe..a8c0f52 100644
--- a/openstackquery/query_blocks/result.py
+++ b/openstackquery/query_blocks/result.py
@@ -36,7 +36,7 @@ def get_prop(self, prop: PropEnum) -> PropValue:
         """
         try:
             return self._prop_enum_cls.get_prop_mapping(prop)(self._obj_result)
-        except AttributeError:
+        except (AttributeError, KeyError):
             return self._default_prop_value
 
     def update_forwarded_properties(self, forwarded_props: Dict[str, PropValue]):