addressed reviews

Author: alexandraBara

nodescraper/plugins/inband/device_enumeration/analyzer_args.py

@@ -23,17 +23,36 @@
 # SOFTWARE.
 #
 ###############################################################################
-from typing import Optional, Union
+from typing import Any, Optional
+
+from pydantic import field_validator
 
 from nodescraper.models import AnalyzerArgs
 
 from .deviceenumdata import DeviceEnumerationDataModel
 
 
 class DeviceEnumerationAnalyzerArgs(AnalyzerArgs):
-    cpu_count: Optional[Union[list[int], int]] = None
-    gpu_count: Optional[Union[list[int], int]] = None
-    vf_count: Optional[Union[list[int], int]] = None
+    cpu_count: Optional[list[int]] = None
+    gpu_count: Optional[list[int]] = None
+    vf_count: Optional[list[int]] = None
+
+    @field_validator("cpu_count", "gpu_count", "vf_count", mode="before")
+    @classmethod
+    def normalize_to_list(cls, v: Any) -> Optional[list[int]]:
+        """Convert single integer values to lists for consistent handling.
+
+        Args:
+            v: The input value (can be int, list[int], or None).
+
+        Returns:
+            Optional[list[int]]: The normalized list value or None.
+        """
+        if v is None:
+            return None
+        if isinstance(v, int):
+            return [v]
+        return v
 
     @classmethod
     def build_from_model(
@@ -48,7 +67,7 @@ def build_from_model(
             DeviceEnumerationAnalyzerArgs: instance of analyzer args class
         """
         return cls(
-            cpu_count=datamodel.cpu_count,
-            gpu_count=datamodel.gpu_count,
-            vf_count=datamodel.vf_count,
+            cpu_count=[datamodel.cpu_count] if datamodel.cpu_count is not None else None,
+            gpu_count=[datamodel.gpu_count] if datamodel.gpu_count is not None else None,
+            vf_count=[datamodel.vf_count] if datamodel.vf_count is not None else None,
         )

nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py

@@ -52,18 +52,13 @@ def analyze_data(
             )
             return self.result
 
-        # Convert to lists if integers, otherwise use as-is
-        cpu_count = [args.cpu_count] if isinstance(args.cpu_count, int) else args.cpu_count
-        gpu_count = [args.gpu_count] if isinstance(args.gpu_count, int) else args.gpu_count
-        vf_count = [args.vf_count] if isinstance(args.vf_count, int) else args.vf_count
-
         checks = {}
-        if cpu_count is not None and cpu_count != []:
-            checks["cpu_count"] = cpu_count
-        if gpu_count is not None and gpu_count != []:
-            checks["gpu_count"] = gpu_count
-        if vf_count is not None and vf_count != []:
-            checks["vf_count"] = vf_count
+        if args.cpu_count is not None and args.cpu_count != []:
+            checks["cpu_count"] = args.cpu_count
+        if args.gpu_count is not None and args.gpu_count != []:
+            checks["gpu_count"] = args.gpu_count
+        if args.vf_count is not None and args.vf_count != []:
+            checks["vf_count"] = args.vf_count
 
         self.result.message = ""
         for check, accepted_counts in checks.items():

nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py

@@ -74,7 +74,6 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
         On Linux, use lscpu and lspci
         On Windows, use WMI and hyper-v cmdlets
         """
-        device_enum = None
         if self.system_info.os_family == OSFamily.LINUX:
             # Count CPU sockets
             cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_LINUX)
@@ -91,31 +90,29 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
             cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_WINDOWS)
             gpu_count_res = self._run_sut_cmd(self.CMD_GPU_COUNT_WINDOWS)
             vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_WINDOWS)
-        cpu_count, gpu_count, vf_count = [None, None, None]
+
+        device_enum = DeviceEnumerationDataModel()
 
         if cpu_count_res.exit_code == 0:
-            cpu_count = int(cpu_count_res.stdout)
+            device_enum.cpu_count = int(cpu_count_res.stdout)
         else:
             self._warning(description="Cannot determine CPU count", command=cpu_count_res)
 
         if gpu_count_res.exit_code == 0:
-            gpu_count = int(gpu_count_res.stdout)
+            device_enum.gpu_count = int(gpu_count_res.stdout)
         else:
             self._warning(description="Cannot determine GPU count", command=gpu_count_res)
 
         if vf_count_res.exit_code == 0:
-            vf_count = int(vf_count_res.stdout)
+            device_enum.vf_count = int(vf_count_res.stdout)
         else:
             self._warning(
                 description="Cannot determine VF count",
                 command=vf_count_res,
                 category=EventCategory.SW_DRIVER,
             )
 
-        if cpu_count or gpu_count or vf_count:
-            device_enum = DeviceEnumerationDataModel(
-                cpu_count=cpu_count, gpu_count=gpu_count, vf_count=vf_count
-            )
+        if device_enum.cpu_count or device_enum.gpu_count or device_enum.vf_count:
             self._log_event(
                 category=EventCategory.PLATFORM,
                 description=f"Counted {device_enum.cpu_count} CPUs, {device_enum.gpu_count} GPUs, {device_enum.vf_count} VFs",
@@ -124,6 +121,7 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
             )
             self.result.message = f"Device Enumeration: {device_enum.model_dump(exclude_none=True)}"
             self.result.status = ExecutionStatus.OK
+            return self.result, device_enum
         else:
             self.result.message = "Device Enumeration info not found"
             self.result.status = ExecutionStatus.EXECUTION_FAILURE
@@ -132,5 +130,4 @@ def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumeratio
                 description=self.result.message,
                 priority=EventPriority.CRITICAL,
             )
-
-        return self.result, device_enum
+            return self.result, None