added test

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("/")])