utest + picking up info through lspci

Author: alexandraBara

nodescraper/plugins/inband/device_enumeration/__init__.py

@@ -23,3 +23,7 @@
 # SOFTWARE.
 #
 ###############################################################################
+from .analyzer_args import DeviceEnumerationAnalyzerArgs
+from .device_enumeration_plugin import DeviceEnumerationPlugin
+
+__all__ = ["DeviceEnumerationPlugin", "DeviceEnumerationAnalyzerArgs"]

nodescraper/plugins/inband/device_enumeration/analyzer_args.py

@@ -23,17 +23,17 @@
 # SOFTWARE.
 #
 ###############################################################################
+from typing import Optional, Union
 
 from nodescraper.models import AnalyzerArgs
-from nodescraper.plugins.inband.device_enumeration.deviceenumdata import (
-    DeviceEnumerationDataModel,
-)
+
+from .deviceenumdata import DeviceEnumerationDataModel
 
 
 class DeviceEnumerationAnalyzerArgs(AnalyzerArgs):
-    cpu_count: list | int = (None,)
-    gpu_count: list | int = (None,)
-    vf_count: list | int = (None,)
+    cpu_count: Optional[Union[list[int], int]] = None
+    gpu_count: Optional[Union[list[int], int]] = None
+    vf_count: Optional[Union[list[int], int]] = None
 
     @classmethod
     def build_from_model(
@@ -47,4 +47,8 @@ def build_from_model(
         Returns:
             DeviceEnumerationAnalyzerArgs: instance of analyzer args class
         """
-        return cls(exp_bios_version=datamodel.bios_version)
+        return cls(
+            cpu_count=datamodel.cpu_count,
+            gpu_count=datamodel.gpu_count,
+            vf_count=datamodel.vf_count,
+        )

nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py

@@ -22,13 +22,15 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 #
-###########################
-
-from errorscraper.event import EventCategory, EventPriority
+###############################################################################
+from typing import Optional
 
+from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus
 from nodescraper.interfaces import DataAnalyzer
-from nodescraper.models import TaskResult, TaskStatus
-from nodescraper.plugins.inband.deviceenumdata import DeviceEnumerationDataModel
+from nodescraper.models import TaskResult
+
+from .analyzer_args import DeviceEnumerationAnalyzerArgs
+from .deviceenumdata import DeviceEnumerationDataModel
 
 
 class DeviceEnumerationAnalyzer(
@@ -45,22 +47,20 @@ def analyze_data(
 
         if not args:
             self.result.message = "Expected Device Enumeration expected data not provided"
-            self.result.status = TaskStatus.NOT_RAN
+            self.result.status = ExecutionStatus.NOT_RAN
             return self.result
 
-        if isinstance(args.cpu_count, int):
-            cpu_count = [args.cpu_count]
-        if isinstance(args.gpu_count, int):
-            gpu_count = [args.gpu_count]
-        if isinstance(args.vf_count, int):
-            vf_count = [args.vf_count]
+        # 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 not in [None, []]:
+        if cpu_count is not None and cpu_count != []:
             checks["cpu_count"] = cpu_count
-        if gpu_count not in [None, []]:
+        if gpu_count is not None and gpu_count != []:
             checks["gpu_count"] = gpu_count
-        if vf_count not in [None, []]:
+        if vf_count is not None and vf_count != []:
             checks["vf_count"] = vf_count
 
         self.result.message = ""
@@ -69,7 +69,7 @@ def analyze_data(
             if actual_count not in accepted_counts:
                 message = f"Expected {check} in {accepted_counts}, but got {actual_count}. "
                 self.result.message += message
-                self.result.status = TaskStatus.ERRORS_DETECTED
+                self.result.status = ExecutionStatus.ERROR
                 self._log_event(
                     category=EventCategory.PLATFORM,
                     description=message,
@@ -78,7 +78,7 @@ def analyze_data(
                     console_log=True,
                 )
         if self.result.message == "":
-            self.result.status = TaskStatus.OK
+            self.result.status = ExecutionStatus.OK
             self.result.message = f"Device Enumeration validated on {checks.keys()}."
 
         return self.result

nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py

@@ -23,20 +23,35 @@
 # SOFTWARE.
 #
 ###############################################################################
-
+from typing import Optional
 
 from nodescraper.base import InBandDataCollector
 from nodescraper.connection.inband.inband import CommandArtifact
-from nodescraper.enums import EventCategory, EventPriority, OSFamily
-from nodescraper.models import TaskResult, TaskStatus
-from nodescraper.plugins.inband.devenumdata import DeviceEnumerationDataModel
+from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus, OSFamily
+from nodescraper.models import TaskResult
+
+from .deviceenumdata import DeviceEnumerationDataModel
 
 
-class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel]):
+class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel, None]):
     """Collect CPU and GPU count"""
 
     DATA_MODEL = DeviceEnumerationDataModel
 
+    # Linux commands
+    CMD_CPU_COUNT_LINUX = "lscpu | grep Socket | awk '{ print $2 }'"
+    CMD_GPU_COUNT_LINUX = "lspci -d 1002: | grep -i 'VGA\\|Display\\|3D' | wc -l"
+    CMD_VF_COUNT_LINUX = "lspci -d 1002: | grep -i 'Virtual Function' | wc -l"
+
+    # Windows commands
+    CMD_CPU_COUNT_WINDOWS = (
+        'powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"'
+    )
+    CMD_GPU_COUNT_WINDOWS = 'powershell -Command "(wmic path win32_VideoController get name | findstr AMD | Measure-Object).Count"'
+    CMD_VF_COUNT_WINDOWS = (
+        'powershell -Command "(Get-VMHostPartitionableGpu | Measure-Object).Count"'
+    )
+
     def _warning(
         self,
         description: str,
@@ -55,30 +70,26 @@ def _warning(
             priority=EventPriority.WARNING,
         )
 
-    def collect_data(self, args=None) -> tuple[TaskResult, DeviceEnumerationDataModel | None]:
+    def collect_data(self, args=None) -> tuple[TaskResult, Optional[DeviceEnumerationDataModel]]:
         """
         Read CPU and GPU count
         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:
-            baremetal_device_id = self.system_info.devid_ep
-            sriov_device_id = self.system_info.devid_ep_vf
+            # Count CPU sockets
+            cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_LINUX)
 
-            cpu_count_res = self._run_system_command("lscpu | grep Socket | awk '{ print $2 }'")
-            gpu_count_res = self._run_system_command(f"lspci -d :{baremetal_device_id} | wc -l")
-            vf_count_res = self._run_system_command(f"lspci -d :{sriov_device_id} | wc -l")
+            # Count all AMD GPUs (vendor ID 1002)
+            gpu_count_res = self._run_sut_cmd(self.CMD_GPU_COUNT_LINUX)
+
+            # Count AMD Virtual Functions
+            vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_LINUX)
         else:
-            cpu_count_res = self._run_system_command(
-                'powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"'
-            )
-            gpu_count_res = self._run_system_command(
-                'powershell -Command "(wmic path win32_VideoController get name | findstr AMD | Measure-Object).Count"'
-            )
-            vf_count_res = self._run_system_command(
-                'powershell -Command "(Get-VMHostPartitionableGpu | Measure-Object).Count"'
-            )
+            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]
 
         if cpu_count_res.exit_code == 0:
@@ -111,10 +122,10 @@ def collect_data(self, args=None) -> tuple[TaskResult, DeviceEnumerationDataMode
                 priority=EventPriority.INFO,
             )
             self.result.message = f"Device Enumeration: {device_enum.model_dump(exclude_none=True)}"
-            self.result.status = TaskStatus.OK
+            self.result.status = ExecutionStatus.OK
         else:
             self.result.message = "Device Enumeration info not found"
-            self.result.status = TaskStatus.EXECUTION_FAILURE
+            self.result.status = ExecutionStatus.EXECUTION_FAILURE
             self._log_event(
                 category=EventCategory.SW_DRIVER,
                 description=self.result.message,

nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py

@@ -28,7 +28,7 @@
 from .analyzer_args import DeviceEnumerationAnalyzerArgs
 from .device_enumeration_analyzer import DeviceEnumerationAnalyzer
 from .device_enumeration_collector import DeviceEnumerationCollector
-from .device_enumerationdata import DeviceEnumerationDataModel
+from .deviceenumdata import DeviceEnumerationDataModel
 
 
 class DeviceEnumerationPlugin(

nodescraper/plugins/inband/device_enumeration/deviceenumdata.py

@@ -23,10 +23,12 @@
 # SOFTWARE.
 #
 ###############################################################################
+from typing import Optional
+
 from nodescraper.models import DataModel
 
 
 class DeviceEnumerationDataModel(DataModel):
-    cpu_count: int | None = None
-    gpu_count: int | None = None
-    vf_count: int | None = None
+    cpu_count: Optional[int] = None
+    gpu_count: Optional[int] = None
+    vf_count: Optional[int] = None

test/unit/plugin/test_device_enumeration_analyzer.py

@@ -0,0 +1,127 @@
+###############################################################################
+#
+# MIT License
+#
+# Copyright (c) 2025 Advanced Micro Devices, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+###############################################################################
+import pytest
+
+from nodescraper.enums.eventcategory import EventCategory
+from nodescraper.enums.eventpriority import EventPriority
+from nodescraper.enums.executionstatus import ExecutionStatus
+from nodescraper.models.systeminfo import OSFamily
+from nodescraper.plugins.inband.device_enumeration.analyzer_args import (
+    DeviceEnumerationAnalyzerArgs,
+)
+from nodescraper.plugins.inband.device_enumeration.device_enumeration_analyzer import (
+    DeviceEnumerationAnalyzer,
+)
+from nodescraper.plugins.inband.device_enumeration.deviceenumdata import (
+    DeviceEnumerationDataModel,
+)
+
+
+@pytest.fixture
+def device_enumeration_analyzer(system_info):
+    return DeviceEnumerationAnalyzer(system_info=system_info)
+
+
+@pytest.fixture
+def device_enumeration_data():
+    return DeviceEnumerationDataModel(cpu_count=4, gpu_count=4, vf_count=8)
+
+
+def test_analyze_passing_linux(system_info, device_enumeration_analyzer, device_enumeration_data):
+    """Test a normal passing case with matching config"""
+    system_info.os_family = OSFamily.LINUX
+
+    args = DeviceEnumerationAnalyzerArgs(cpu_count=4, gpu_count=4, vf_count=8)
+
+    result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=args)
+
+    assert result.status == ExecutionStatus.OK
+    assert len(result.events) == 0
+
+
+def test_analyze_passing_windows(system_info, device_enumeration_analyzer, device_enumeration_data):
+    """Test a normal passing case on Windows"""
+    system_info.os_family = OSFamily.WINDOWS
+
+    args = DeviceEnumerationAnalyzerArgs(gpu_count=4, vf_count=8)
+
+    result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=args)
+
+    assert result.status == ExecutionStatus.OK
+    assert len(result.events) == 0
+
+
+def test_analyze_no_args(device_enumeration_analyzer, device_enumeration_data):
+    """Test with no analyzer args provided"""
+
+    result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=None)
+
+    assert result.status == ExecutionStatus.NOT_RAN
+    assert "Expected Device Enumeration expected data not provided" in result.message
+    assert len(result.events) == 0
+
+
+def test_analyze_unexpected_counts(device_enumeration_analyzer, device_enumeration_data):
+    """Test with config specifying different device counts"""
+
+    args = DeviceEnumerationAnalyzerArgs(cpu_count=1, gpu_count=10)
+
+    result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=args)
+
+    assert result.status == ExecutionStatus.ERROR
+    assert "but got" in result.message
+
+    for event in result.events:
+        assert event.priority == EventPriority.CRITICAL
+        assert event.category == EventCategory.PLATFORM.value
+
+
+def test_analyze_mismatched_cpu_count(device_enumeration_analyzer):
+    """Test with invalid device enumeration on SUT"""
+
+    data = DeviceEnumerationDataModel(cpu_count=5, gpu_count=4, vf_count=8)
+    args = DeviceEnumerationAnalyzerArgs(cpu_count=4, gpu_count=4)
+
+    result = device_enumeration_analyzer.analyze_data(data=data, args=args)
+
+    assert result.status == ExecutionStatus.ERROR
+    assert "but got" in result.message
+
+    for event in result.events:
+        assert event.priority == EventPriority.CRITICAL
+        assert event.category == EventCategory.PLATFORM.value
+
+
+def test_analyze_list_of_accepted_counts(device_enumeration_analyzer):
+    """Test with a list of acceptable counts"""
+
+    data = DeviceEnumerationDataModel(cpu_count=4, gpu_count=4, vf_count=8)
+    args = DeviceEnumerationAnalyzerArgs(cpu_count=[2, 4, 8], gpu_count=[4, 8])
+
+    result = device_enumeration_analyzer.analyze_data(data=data, args=args)
+
+    assert result.status == ExecutionStatus.OK
+    assert len(result.events) == 0