Add test for code when disk quota is reached

Signed-off-by: Fabrice Normandin <[email protected]>

tests/integration/test_code.py

@@ -8,14 +8,17 @@
 import time
 from datetime import timedelta
 from logging import getLogger as get_logger
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 from pytest_regressions.file_regression import FileRegressionFixture
 
 from milatools.cli.code import code
 from milatools.cli.commands import code_v1
-from milatools.cli.utils import get_hostname_to_use_for_compute_node
+from milatools.cli.utils import MilatoolsUserError, get_hostname_to_use_for_compute_node
+from milatools.utils import disk_quota
 from milatools.utils.compute_node import ComputeNode
+from milatools.utils.disk_quota import check_disk_quota, check_disk_quota_v1
 from milatools.utils.remote_v1 import RemoteV1
 from milatools.utils.remote_v2 import RemoteV2
 
@@ -183,6 +186,74 @@ def filter_line(line: str) -> str:
     file_regression.check(filter_captured_output(captured_output))
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_v1", [False, True], ids=["v2", "v1"])
+async def test_code_with_disk_quota_reached(
+    monkeypatch: pytest.MonkeyPatch, use_v1: bool
+):
+    if use_v1:
+        from milatools.cli import commands
+
+        # Makes the test slightly quicker to run.
+        monkeypatch.setattr(commands, RemoteV1.__name__, Mock(spec=RemoteV1))
+
+        def _mock_check_disk_quota_v1(remote: RemoteV1 | RemoteV2):
+            raise MilatoolsUserError(
+                "ERROR: Your disk quota on the $HOME filesystem is exceeded! "
+            )
+
+        mock_check_disk_quota = Mock(
+            spec=check_disk_quota_v1, side_effect=_mock_check_disk_quota_v1
+        )
+        monkeypatch.setattr(
+            disk_quota, check_disk_quota_v1.__name__, mock_check_disk_quota
+        )
+        monkeypatch.setattr(
+            commands, check_disk_quota_v1.__name__, mock_check_disk_quota
+        )
+        with pytest.raises(MilatoolsUserError):
+            code_v1(
+                path="bob",
+                command="echo",  # replace the usual `code` with `echo` for testing.
+                persist=False,
+                job=None,
+                node="bobobo",  # to avoid accidental sallocs.
+                alloc=[],
+                cluster="bob",
+            )
+        mock_check_disk_quota.assert_called_once()
+    else:
+        from milatools.cli import code
+
+        # Makes the test quicker to run by avoiding connecting to the cluster.
+        monkeypatch.setattr(code, RemoteV2.__name__, Mock(spec=RemoteV2))
+
+        async def _mock_check_disk_quota(remote: RemoteV1 | RemoteV2):
+            raise MilatoolsUserError(
+                "ERROR: Your disk quota on the $HOME filesystem is exceeded! "
+            )
+
+        mock_check_disk_quota = AsyncMock(
+            spec=check_disk_quota, side_effect=_mock_check_disk_quota
+        )
+        monkeypatch.setattr(
+            disk_quota, check_disk_quota.__name__, mock_check_disk_quota
+        )
+
+        monkeypatch.setattr(code, check_disk_quota.__name__, mock_check_disk_quota)
+        with pytest.raises(MilatoolsUserError):
+            await code.code(
+                path="bob",
+                command="echo",  # replace the usual `code` with `echo` for testing.
+                persist=False,
+                job=None,
+                node="bobobo",  # to avoid accidental sallocs.
+                alloc=[],
+                cluster="bob",
+            )
+        mock_check_disk_quota.assert_called_once()
+
+
 @pytest.mark.slow
 @launches_jobs
 @PARAMIKO_SSH_BANNER_BUG