diff --git a/openfe/tests/utils/test_system_probe.py b/openfe/tests/utils/test_system_probe.py index 386e8adc5..0b274fab8 100644 --- a/openfe/tests/utils/test_system_probe.py +++ b/openfe/tests/utils/test_system_probe.py @@ -1,9 +1,11 @@ import contextlib from collections import namedtuple +import pathlib import sys from unittest.mock import Mock, patch import psutil +from psutil._common import sdiskusage import pytest from openfe.utils.system_probe import ( @@ -38,6 +40,17 @@ ) +def fake_disk_usage(path): + if str(path) == "/foo": + return sdiskusage( + total=1958854045696, used=1232985415680, free=626288726016, percent=66.3 + ) + if str(path) == "/bar": + return sdiskusage( + total=4000770252800, used=1678226952192, free=2322615496704, percent=41.9 + ) + + @contextlib.contextmanager def patch_system(): # single patch to fully patch the system @@ -91,6 +104,9 @@ def patch_system(): ) ), ) + patch_psutil_disk_usage = patch( + "psutil.disk_usage", Mock(side_effect=fake_disk_usage) + ) # assumes that each shell command is called in only one way cmd_to_output = { @@ -111,10 +127,11 @@ def patch_system(): ) with contextlib.ExitStack() as stack: for ctx in [ + patch_check_output, patch_hostname, patch_psutil_Process_as_dict, - patch_check_output, patch_psutil_Process_rlimit, + patch_psutil_disk_usage, patch_psutil_virtual_memory, ]: stack.enter_context(ctx) @@ -229,6 +246,30 @@ def test_get_disk_usage(): assert disk_info == expected_disk_info +@pytest.mark.skipif( + sys.platform == "darwin", reason="test requires psutil.Process.rlimit" +) +def test_get_disk_usage_with_path(): + with patch_system(): + disk_info = _get_disk_usage(paths=[pathlib.Path("/foo"), pathlib.Path("/bar")]) + expected_disk_info = { + "/bar": { + "available": "2.1T", + "percent_used": "42%", + "size": "3.6T", + "used": "1.5T", + }, + "/foo": { + "available": "583.3G", + "percent_used": "66%", + "size": "1.8T", + "used": "1.1T", + }, + } + + assert disk_info == expected_disk_info + + @pytest.mark.skipif( sys.platform == "darwin", reason="test requires psutil.Process.rlimit" ) @@ -317,3 +358,4 @@ def test_probe_system(): def test_probe_system_smoke_test(): _probe_system() + _probe_system(paths=[pathlib.Path("/")])