diff --git a/agentops/host_env.py b/agentops/host_env.py index 2c298b82..42666ab7 100644 --- a/agentops/host_env.py +++ b/agentops/host_env.py @@ -2,6 +2,7 @@ import psutil import socket from .helpers import get_agentops_version +from .log_config import logger import importlib.metadata import os import sys @@ -115,14 +116,19 @@ def get_disk_details(): partitions = psutil.disk_partitions() disk_info = {} for partition in partitions: - usage = psutil.disk_usage(partition.mountpoint) - disk_info[partition.device] = { - "Mountpoint": partition.mountpoint, - "Total": f"{usage.total / (1024**3):.2f} GB", - "Used": f"{usage.used / (1024**3):.2f} GB", - "Free": f"{usage.free / (1024**3):.2f} GB", - "Percentage": f"{usage.percent}%", - } + try: + usage = psutil.disk_usage(partition.mountpoint) + disk_info[partition.device] = { + "Mountpoint": partition.mountpoint, + "Total": f"{usage.total / (1024**3):.2f} GB", + "Used": f"{usage.used / (1024**3):.2f} GB", + "Free": f"{usage.free / (1024**3):.2f} GB", + "Percentage": f"{usage.percent}%", + } + except OSError as inaccessible: + # Skip inaccessible partitions, such as removable drives with no media + logger.debug("Mountpoint %s inaccessible: %s", partition.mountpoint, inaccessible) + return disk_info diff --git a/tests/test_host_env.py b/tests/test_host_env.py new file mode 100644 index 00000000..8da16394 --- /dev/null +++ b/tests/test_host_env.py @@ -0,0 +1,43 @@ +from unittest.mock import patch +from agentops import host_env + +# noinspection PyProtectedMember +from psutil._common import sdiskpart, sdiskusage + + +def mock_partitions(): + return [ + sdiskpart(device='/dev/sda1', mountpoint='/', fstype='ext4', + opts='rw,relatime', maxfile=0, maxpath=0), + sdiskpart(device='z:\\', mountpoint='z:\\', fstype='ntfs', + opts='rw,relatime', maxfile=0, maxpath=0), + ] + + +def mock_disk_usage(partition): + if partition == '/': + return sdiskusage(total=(1024**3), used=0, free=(1024**3), percent=100) + else: + raise PermissionError("Device access exception should have been caught") + + +class TestHostEnv: + @patch('psutil.disk_partitions', new=lambda: [mock_partitions()[0]]) + @patch('psutil.disk_usage', new=mock_disk_usage) + def test_disk_info(self): + self.assert_disk_info() + + @patch('psutil.disk_partitions', new=mock_partitions) + @patch('psutil.disk_usage', new=mock_disk_usage) + def test_disk_info_skips_oserror(self): + self.assert_disk_info() + + def assert_disk_info(self): + disk_info = host_env.get_disk_details() + assert list(disk_info.keys()) == ["/dev/sda1"] + sda1 = disk_info["/dev/sda1"] + assert sda1["Mountpoint"] == "/" + assert sda1["Total"] == "1.00 GB" + assert sda1["Used"] == "0.00 GB" + assert sda1["Free"] == "1.00 GB" + assert sda1["Percentage"] == "100%"