Add tests for minimum free disk space check

BUCK

@@ -20,3 +20,14 @@ genrule(
     out = 'dist-upgrader',
     cmd = 'cp "$(location :dist-upgrader.pex)" "$OUT" && chmod +x "$OUT"',
 )
+
+python_test(
+    name = 'dist-upgrader.tests',
+    srcs = glob(['./tests/*.py']),
+    deps = [
+        '//pleskdistup:lib',
+    ],
+    base_module = '',
+    main_module = 'tests.test_main',
+    platform = 'py3',
+)

pleskdistup/actions/common_checks.py

@@ -158,6 +158,60 @@ class MinFreeDiskSpaceViolation(typing.NamedTuple):
     """Paths belonging to this device."""
 
 
+def check_min_free_disk_space(requirements: typing.Dict[str, int]) -> typing.List[MinFreeDiskSpaceViolation]:
+    """Check minimum free disk space according to the requirements.
+
+    Args:
+        requirements: A dictionary mapping paths to minimum free disk
+            space (in bytes) on the devices containing them.
+
+    Returns:
+        A list of filesystems with insufficient free disk space.
+    """
+    cmd = [
+        "/bin/findmnt", "--output", "source,target,avail",
+        "--bytes", "--json", "-T",
+    ]
+    violations: typing.List[MinFreeDiskSpaceViolation] = []
+    filesystems: typing.Dict[str, dict] = {}
+    for path, req in requirements.items():
+        log.debug(f"Checking {path!r} minimum free disk space requirement of {req}")
+        proc = subprocess.run(
+            cmd + [path],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            check=True,
+            universal_newlines=True,
+        )
+        log.debug(
+            f"Command {cmd + [path]} returned {proc.returncode}, "
+            f"stdout: '{proc.stdout}', stderr: '{proc.stderr}'"
+        )
+        fs_data = json.loads(proc.stdout)["filesystems"][0]
+        if fs_data["source"] not in filesystems:
+            log.debug(f"Discovered new filesystem {fs_data}")
+            fs_data["req"] = 0
+            fs_data["paths"] = set()
+            filesystems[fs_data["source"]] = fs_data
+        log.debug(
+            f"Adding space requirement of {req} to "
+            f"{filesystems[fs_data['source']]}"
+        )
+        filesystems[fs_data["source"]]["req"] += req
+        filesystems[fs_data["source"]]["paths"].add(path)
+    for dev, fs_data in filesystems.items():
+        if fs_data["req"] > fs_data["avail"]:
+            violations.append(
+                MinFreeDiskSpaceViolation(
+                    dev,
+                    fs_data["req"],
+                    fs_data["avail"],
+                    fs_data["paths"],
+                )
+            )
+    return violations
+
+
 class AssertMinFreeDiskSpace(action.CheckAction):
     """Check if there's enough free disk space.
 
@@ -167,13 +221,14 @@ class AssertMinFreeDiskSpace(action.CheckAction):
         name: Name of the check.
     """
     violations: typing.List[MinFreeDiskSpaceViolation]
-    """List of filesystems with insiffucient free disk space."""
+    """List of filesystems with insufficient free disk space."""
 
     def __init__(
         self,
         requirements: typing.Dict[str, int],
         name: str = "check if there's enough free disk space",
     ):
+        super().__init__()
         self.requirements = requirements
         self.name = name
         self.violations = []
@@ -195,46 +250,6 @@ def _update_description(self) -> None:
     def _do_check(self) -> bool:
         """Perform the check."""
         log.debug("Checking minimum free disk space")
-        cmd = [
-            "/bin/findmnt", "--output", "source,target,avail",
-            "--bytes", "--json", "-T",
-        ]
-        self.violations = []
-        filesystems: typing.Dict[str, dict] = {}
-        for path, req in self.requirements.items():
-            log.debug(f"Checking {path!r} minimum free disk space requirement of {req}")
-            proc = subprocess.run(
-                cmd + [path],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                check=True,
-                universal_newlines=True,
-            )
-            log.debug(
-                f"Command {cmd + [path]} returned {proc.returncode}, "
-                f"stdout: '{proc.stdout}', stderr: '{proc.stderr}'"
-            )
-            fs_data = json.loads(proc.stdout)["filesystems"][0]
-            if fs_data["source"] not in filesystems:
-                log.debug(f"Discovered new filesystem {fs_data}")
-                fs_data["req"] = 0
-                fs_data["paths"] = set()
-                filesystems[fs_data["source"]] = fs_data
-            log.debug(
-                f"Adding space requirement of {req} to "
-                f"{filesystems[fs_data['source']]}"
-            )
-            filesystems[fs_data["source"]]["req"] += req
-            filesystems[fs_data["source"]]["paths"].add(path)
-        for dev, fs_data in filesystems.items():
-            if fs_data["req"] > fs_data["avail"]:
-                self.violations.append(
-                    MinFreeDiskSpaceViolation(
-                        dev,
-                        fs_data["req"],
-                        fs_data["avail"],
-                        fs_data["paths"],
-                    )
-                )
+        self.violations = check_min_free_disk_space(self.requirements)
         self._update_description()
         return len(self.violations) == 0

tests/common_checks.py

@@ -0,0 +1,188 @@
+# Copyright 2023-2024. WebPros International GmbH. All rights reserved.
+
+import subprocess
+import unittest
+from unittest import mock
+
+from pleskdistup.actions.common_checks import (
+    AssertMinFreeDiskSpace,
+    MinFreeDiskSpaceViolation,
+    check_min_free_disk_space,
+)
+
+
+class FindmntMock:
+    def __init__(self, args, results):
+        self.args = args
+        self.results = results
+
+    def __call__(self, cmd, *args, **kwargs):
+        target = cmd[-1]
+        out = """
+{{
+   "filesystems": [
+      {{
+         "source": "{source}",
+         "target": "{target}",
+         "avail": {avail}
+      }}
+   ]
+}}
+""".format(**self.results[target])
+        return subprocess.CompletedProcess(
+            args=self.args,
+            returncode=0,
+            stdout=out,
+            stderr="",
+        )
+
+
+class TestCheckMinFreeDiskSpace(unittest.TestCase):
+    _requirements = {
+        # Space requirements in bytes
+        "/boot": 150 * 1024**2,
+        "/opt": 150 * 1024**2,
+        "/usr": 1800 * 1024**2,
+        "/var": 2000 * 1024**2,
+    }
+    _findmnt_cmd = [
+        "/bin/findmnt", "--output", "source,target,avail",
+        "--bytes", "--json", "-T",
+    ]
+
+    def _get_expected_calls(self, paths):
+        return [
+            mock.call(
+                self._findmnt_cmd + [path],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                check=True,
+                universal_newlines=True,
+            ) for path in paths
+        ]
+
+    def setUp(self):
+        self.maxDiff = None
+
+    def test_no_violations(self):
+        req_sum = sum(
+            self._requirements[k] for k in ("/boot", "/opt", "/usr")
+        )
+        findmnt_results = {
+            "/boot": {"source": "/dev/sda1", "target": "/", "avail": req_sum},
+            "/opt": {"source": "/dev/sda1", "target": "/", "avail": req_sum},
+            "/usr": {"source": "/dev/sda1", "target": "/", "avail": req_sum},
+            "/var": {"source": "/dev/sdb1", "target": "/var", "avail": self._requirements["/var"]},
+        }
+        with mock.patch(
+            'subprocess.run',
+            side_effect=FindmntMock(self._findmnt_cmd, findmnt_results),
+        ) as run_mock:
+            violations = check_min_free_disk_space(self._requirements)
+            self.assertEqual(violations, [])
+            self.assertListEqual(
+                run_mock.mock_calls,
+                self._get_expected_calls(self._requirements),
+            )
+
+    def test_violations(self):
+        # Insufficient space on /dev/sda1, /dev/sda2
+        dev_avail = {
+            "/dev/sda1": self._requirements["/boot"] + self._requirements["/opt"] - 1,
+            "/dev/sda2": self._requirements["/usr"] - 1,
+            "/dev/sda3": self._requirements["/var"],
+        }
+        findmnt_results = {
+            "/boot": {"source": "/dev/sda1", "target": "/"},
+            "/opt": {"source": "/dev/sda1", "target": "/"},
+            "/usr": {"source": "/dev/sda2", "target": "/usr"},
+            "/var": {"source": "/dev/sda3", "target": "/var"},
+        }
+        for path, fs_data in findmnt_results.items():
+            fs_data["avail"] = dev_avail[fs_data["source"]]
+        expected_violations = [
+            MinFreeDiskSpaceViolation(
+                "/dev/sda1",
+                self._requirements["/boot"] + self._requirements["/opt"],
+                dev_avail["/dev/sda1"],
+                {"/boot", "/opt"}
+            ),
+            MinFreeDiskSpaceViolation(
+                "/dev/sda2",
+                self._requirements["/usr"],
+                dev_avail["/dev/sda2"],
+                {"/usr"}
+            ),
+        ]
+        with mock.patch(
+            'subprocess.run',
+            side_effect=FindmntMock(self._findmnt_cmd, findmnt_results),
+        ) as run_mock:
+            violations = check_min_free_disk_space(self._requirements)
+            self.assertListEqual(violations, expected_violations)
+            self.assertListEqual(
+                run_mock.mock_calls,
+                self._get_expected_calls(self._requirements),
+            )
+
+
+class TestAssertMinFreeDiskSpace(unittest.TestCase):
+    _requirements = {
+        # Space requirements in bytes
+        "/boot": 150 * 1024**2,
+        "/opt": 150 * 1024**2,
+        "/usr": 1800 * 1024**2,
+        "/var": 2000 * 1024**2,
+    }
+
+    def setUp(self):
+        self.maxDiff = None
+
+    def test_description_pass(self):
+        with mock.patch(
+            'pleskdistup.actions.common_checks.check_min_free_disk_space',
+            return_value=[],
+        ) as check_min_free_disk_space_mock:
+            assrt = AssertMinFreeDiskSpace(self._requirements)
+            self.assertEqual(assrt.description, "")
+            self.assertTrue(assrt.do_check())
+            self.assertEqual(assrt.description, "")
+            check_min_free_disk_space_mock.assert_called_once_with(self._requirements)
+
+    def test_description_fail(self):
+        # Insufficient space on /dev/sda1, /dev/sda2
+        dev_avail = {
+            "/dev/sda1": self._requirements["/boot"] + self._requirements["/opt"] - 1,
+            "/dev/sda2": self._requirements["/usr"] - 1,
+            "/dev/sda3": self._requirements["/var"],
+        }
+        violations = [
+            MinFreeDiskSpaceViolation(
+                "/dev/sda1",
+                self._requirements["/boot"] + self._requirements["/opt"],
+                dev_avail["/dev/sda1"],
+                {"/boot", "/opt"}
+            ),
+            MinFreeDiskSpaceViolation(
+                "/dev/sda2",
+                self._requirements["/usr"],
+                dev_avail["/dev/sda2"],
+                {"/usr"}
+            ),
+        ]
+        expected_description = "There's not enough free disk space: "
+        expected_description += ", ".join(
+            f"on filesystem {v.dev!r} for "
+            f"{', '.join(repr(p) for p in sorted(v.paths))} "
+            f"(need {v.req_bytes / 1024**2} MiB, "
+            f"got {v.avail_bytes / 1024**2} MiB)" for v in violations
+        )
+        with mock.patch(
+            'pleskdistup.actions.common_checks.check_min_free_disk_space',
+            return_value=violations,
+        ) as check_min_free_disk_space_mock:
+            assrt = AssertMinFreeDiskSpace(self._requirements)
+            self.assertEqual(assrt.description, "")
+            self.assertFalse(assrt.do_check())
+            self.assertEqual(assrt.description, expected_description)
+            check_min_free_disk_space_mock.assert_called_once_with(self._requirements)

tests/test_main.py

@@ -0,0 +1,27 @@
+# Copyright 2023-2024. WebPros International GmbH. All rights reserved.
+
+import logging
+import sys
+
+# Import Buck default test main module
+# See https://buck.build/rule/python_test.html
+import __test_main__
+
+from pleskdistup.common import log
+
+
+def init_logger():
+    log.init_logger(
+        ["tests.log"],
+        [],
+        loglevel=logging.DEBUG,
+    )
+
+
+def main():
+    init_logger()
+    __test_main__.main(sys.argv)
+
+
+if __name__ == '__main__':
+    main()