diff --git a/nodescraper/models/systeminfo.py b/nodescraper/models/systeminfo.py index d2be4ae4..e82d6212 100644 --- a/nodescraper/models/systeminfo.py +++ b/nodescraper/models/systeminfo.py @@ -41,3 +41,4 @@ class SystemInfo(BaseModel): platform: Optional[str] = None metadata: Optional[dict] = Field(default_factory=dict) location: Optional[SystemLocation] = SystemLocation.LOCAL + vendorid_ep: int = 0x1002 diff --git a/nodescraper/plugins/inband/device_enumeration/__init__.py b/nodescraper/plugins/inband/device_enumeration/__init__.py new file mode 100644 index 00000000..a5073399 --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/__init__.py @@ -0,0 +1,29 @@ +############################################################################### +# +# 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. +# +############################################################################### +from .analyzer_args import DeviceEnumerationAnalyzerArgs +from .device_enumeration_plugin import DeviceEnumerationPlugin + +__all__ = ["DeviceEnumerationPlugin", "DeviceEnumerationAnalyzerArgs"] diff --git a/nodescraper/plugins/inband/device_enumeration/analyzer_args.py b/nodescraper/plugins/inband/device_enumeration/analyzer_args.py new file mode 100644 index 00000000..8f74ed00 --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/analyzer_args.py @@ -0,0 +1,73 @@ +############################################################################### +# +# 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. +# +############################################################################### +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[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( + cls, datamodel: DeviceEnumerationDataModel + ) -> "DeviceEnumerationAnalyzerArgs": + """build analyzer args from data model + + Args: + datamodel (DeviceEnumerationDataModel): data model for plugin + + Returns: + DeviceEnumerationAnalyzerArgs: instance of analyzer args class + """ + return cls( + 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, + ) diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py new file mode 100644 index 00000000..7cf39335 --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_analyzer.py @@ -0,0 +1,81 @@ +############################################################################### +# +# 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. +# +############################################################################### +from typing import Optional + +from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus +from nodescraper.interfaces import DataAnalyzer +from nodescraper.models import TaskResult + +from .analyzer_args import DeviceEnumerationAnalyzerArgs +from .deviceenumdata import DeviceEnumerationDataModel + + +class DeviceEnumerationAnalyzer( + DataAnalyzer[DeviceEnumerationDataModel, DeviceEnumerationAnalyzerArgs] +): + """Check Device Enumeration matches expected cpu and gpu count + supported by all OSs, SKUs, and platforms.""" + + DATA_MODEL = DeviceEnumerationDataModel + + def analyze_data( + self, data: DeviceEnumerationDataModel, args: Optional[DeviceEnumerationAnalyzerArgs] = None + ) -> TaskResult: + + if args is None: + self.result.status = ExecutionStatus.NOT_RAN + self.result.message = ( + "Expected Device Enumeration data not provided, skipping analysis." + ) + return self.result + + checks = {} + 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(): + actual_count = getattr(data, check) + 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 = ExecutionStatus.ERROR + self._log_event( + category=EventCategory.PLATFORM, + description=message, + data={check: actual_count}, + priority=EventPriority.CRITICAL, + console_log=True, + ) + if self.result.message == "": + self.result.status = ExecutionStatus.OK + self.result.message = f"Device Enumeration validated on {checks.keys()}." + + return self.result diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py new file mode 100644 index 00000000..88506f6e --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_collector.py @@ -0,0 +1,133 @@ +############################################################################### +# +# 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. +# +############################################################################### +from typing import Optional + +from nodescraper.base import InBandDataCollector +from nodescraper.connection.inband.inband import CommandArtifact +from nodescraper.enums import EventCategory, EventPriority, ExecutionStatus, OSFamily +from nodescraper.models import TaskResult + +from .deviceenumdata import DeviceEnumerationDataModel + + +class DeviceEnumerationCollector(InBandDataCollector[DeviceEnumerationDataModel, None]): + """Collect CPU and GPU count""" + + DATA_MODEL = DeviceEnumerationDataModel + + CMD_CPU_COUNT_LINUX = "lscpu | grep Socket | awk '{ print $2 }'" + CMD_GPU_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'VGA\\|Display\\|3D' | wc -l" + CMD_VF_COUNT_LINUX = "lspci -d {vendorid_ep}: | grep -i 'Virtual Function' | wc -l" + + 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, + command: CommandArtifact, + category: EventCategory = EventCategory.PLATFORM, + ): + self._log_event( + category=category, + description=description, + data={ + "command": command.command, + "stdout": command.stdout, + "stderr": command.stderr, + "exit_code": command.exit_code, + }, + priority=EventPriority.WARNING, + ) + + 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 + """ + if self.system_info.os_family == OSFamily.LINUX: + # Count CPU sockets + cpu_count_res = self._run_sut_cmd(self.CMD_CPU_COUNT_LINUX) + + # Count all AMD GPUs + vendor_id = format(self.system_info.vendorid_ep, "x") + gpu_count_res = self._run_sut_cmd( + self.CMD_GPU_COUNT_LINUX.format(vendorid_ep=vendor_id) + ) + + # Count AMD Virtual Functions + vf_count_res = self._run_sut_cmd(self.CMD_VF_COUNT_LINUX.format(vendorid_ep=vendor_id)) + else: + 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) + + device_enum = DeviceEnumerationDataModel() + + if cpu_count_res.exit_code == 0: + 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: + 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: + 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 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", + data=device_enum.model_dump(exclude_none=True), + priority=EventPriority.INFO, + ) + 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 + self._log_event( + category=EventCategory.SW_DRIVER, + description=self.result.message, + priority=EventPriority.CRITICAL, + ) + return self.result, None diff --git a/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py b/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py new file mode 100644 index 00000000..baf2aa2d --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/device_enumeration_plugin.py @@ -0,0 +1,45 @@ +############################################################################### +# +# 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. +# +############################################################################### +from nodescraper.base import InBandDataPlugin + +from .analyzer_args import DeviceEnumerationAnalyzerArgs +from .device_enumeration_analyzer import DeviceEnumerationAnalyzer +from .device_enumeration_collector import DeviceEnumerationCollector +from .deviceenumdata import DeviceEnumerationDataModel + + +class DeviceEnumerationPlugin( + InBandDataPlugin[DeviceEnumerationDataModel, None, DeviceEnumerationAnalyzerArgs] +): + """Plugin for collection and analysis of BIOS data""" + + DATA_MODEL = DeviceEnumerationDataModel + + COLLECTOR = DeviceEnumerationCollector + + ANALYZER = DeviceEnumerationAnalyzer + + ANALYZER_ARGS = DeviceEnumerationAnalyzerArgs diff --git a/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py b/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py new file mode 100644 index 00000000..74939209 --- /dev/null +++ b/nodescraper/plugins/inband/device_enumeration/deviceenumdata.py @@ -0,0 +1,34 @@ +############################################################################### +# +# 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. +# +############################################################################### +from typing import Optional + +from nodescraper.models import DataModel + + +class DeviceEnumerationDataModel(DataModel): + cpu_count: Optional[int] = None + gpu_count: Optional[int] = None + vf_count: Optional[int] = None diff --git a/test/unit/plugin/test_device_enumeration_analyzer.py b/test/unit/plugin/test_device_enumeration_analyzer.py new file mode 100644 index 00000000..c58c9ea0 --- /dev/null +++ b/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 - should skip analysis""" + + result = device_enumeration_analyzer.analyze_data(data=device_enumeration_data, args=None) + + assert result.status == ExecutionStatus.NOT_RAN + assert "Expected Device Enumeration data not provided, skipping analysis." 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 diff --git a/test/unit/plugin/test_device_enumeration_collector.py b/test/unit/plugin/test_device_enumeration_collector.py new file mode 100644 index 00000000..a5a1ef30 --- /dev/null +++ b/test/unit/plugin/test_device_enumeration_collector.py @@ -0,0 +1,143 @@ +############################################################################### +# +# 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. +# +############################################################################### +from unittest.mock import MagicMock + +import pytest + +from nodescraper.enums.executionstatus import ExecutionStatus +from nodescraper.enums.systeminteraction import SystemInteractionLevel +from nodescraper.models.systeminfo import OSFamily +from nodescraper.plugins.inband.device_enumeration.device_enumeration_collector import ( + DeviceEnumerationCollector, +) +from nodescraper.plugins.inband.device_enumeration.deviceenumdata import ( + DeviceEnumerationDataModel, +) + + +@pytest.fixture +def device_enumeration_collector(system_info, conn_mock): + return DeviceEnumerationCollector( + system_info=system_info, + system_interaction_level=SystemInteractionLevel.PASSIVE, + connection=conn_mock, + ) + + +def test_collect_linux(system_info, device_enumeration_collector): + """Test linux typical output""" + system_info.os_family = OSFamily.LINUX + + device_enumeration_collector._run_sut_cmd = MagicMock( + side_effect=[ + MagicMock( + exit_code=0, + stdout="2", + stderr="", + command="lscpu | grep Socket | awk '{ print $2 }'", + ), + MagicMock( + exit_code=0, + stdout="8", + stderr="", + command="lspci -d 1002: | grep -i 'VGA\\|Display\\|3D' | wc -l", + ), + MagicMock( + exit_code=0, + stdout="0", + stderr="", + command="lspci -d 1002: | grep -i 'Virtual Function' | wc -l", + ), + ] + ) + + result, data = device_enumeration_collector.collect_data() + assert result.status == ExecutionStatus.OK + assert data == DeviceEnumerationDataModel(cpu_count=2, gpu_count=8, vf_count=0) + + +def test_collect_windows(system_info, device_enumeration_collector): + """Test windows typical output""" + system_info.os_family = OSFamily.WINDOWS + + device_enumeration_collector._run_sut_cmd = MagicMock( + side_effect=[ + MagicMock( + exit_code=0, + stdout="2", + stderr="", + command='powershell -Command "(Get-WmiObject -Class Win32_Processor | Measure-Object).Count"', + ), + MagicMock( + exit_code=0, + stdout="8", + stderr="", + command='powershell -Command "(wmic path win32_VideoController get name | findstr AMD | Measure-Object).Count"', + ), + MagicMock( + exit_code=0, + stdout="8", + stderr="", + command='powershell -Command "(Get-VMHostPartitionableGpu | Measure-Object).Count"', + ), + ] + ) + + result, data = device_enumeration_collector.collect_data() + assert result.status == ExecutionStatus.OK + assert data == DeviceEnumerationDataModel(cpu_count=2, gpu_count=8, vf_count=8) + + +def test_collect_error(system_info, device_enumeration_collector): + """Test with bad exit code""" + system_info.os_family = OSFamily.LINUX + + device_enumeration_collector._run_sut_cmd = MagicMock( + side_effect=[ + MagicMock( + exit_code=1, + stdout="some output", + stderr="command failed", + command="lscpu | grep Socket | awk '{ print $2 }'", + ), + MagicMock( + exit_code=1, + stdout="some output", + stderr="command failed", + command="lspci -d 1002: | grep -i 'VGA\\|Display\\|3D' | wc -l", + ), + MagicMock( + exit_code=1, + stdout="some output", + stderr="command failed", + command="lspci -d 1002: | grep -i 'Virtual Function' | wc -l", + ), + ] + ) + + result, data = device_enumeration_collector.collect_data() + assert result.status == ExecutionStatus.EXECUTION_FAILURE + assert data is None