fix(disk-usage): correct disk usage calculation for study variants wi…

…th simulation results (#1926)

Merge pull request #1926 from AntaresSimulatorTeam/bugfix/738-variant-disk-usage (ANT-738)

antarest/study/service.py

@@ -127,6 +127,7 @@
 
 
 def get_disk_usage(path: t.Union[str, Path]) -> int:
+    """Calculate the total disk usage (in bytes) of a study in a compressed file or directory."""
     path = Path(path)
     if path.suffix.lower() in {".zip", "7z"}:
         return os.path.getsize(path)
@@ -2358,19 +2359,23 @@ def upgrade_study(
 
     def get_disk_usage(self, uuid: str, params: RequestParameters) -> int:
         """
-        This function computes the disk size used to store the study with
-        id=`uuid` if such study exists and user has permissions
-        otherwise it raises an error
+        Calculates the size of the disk used to store the study if the user has permissions.
+
+        The calculation of disk space concerns the entire study directory.
+        In the case of a variant, the snapshot folder must be taken into account, as well as the outputs.
 
         Args:
-            uuid: the study id
-            params: user request parameters
+            uuid: the study ID.
+            params: user request parameters.
+
+        Returns:
+            Disk usage of the study in bytes.
 
-        return:
-            disk usage of the study with id = `uuid`
+        Raises:
+            UserHasNotPermissionError: If the user does not have the READ permissions (HTTP status 403).
         """
         study = self.get_study(uuid=uuid)
         assert_permission(params.user, study, StudyPermissionType.READ)
-        study_path = self.storage_service.get_storage(study).get_study_path(study)
-        # If the study is a variant, it's possible that it only exists in db and not on disk. If so, we return 0.
-        return get_disk_usage(path=str(study_path)) if study_path.exists() else 0
+        study_path = self.storage_service.raw_study_service.get_study_path(study)
+        # If the study is a variant, it's possible that it only exists in DB and not on disk. If so, we return 0.
+        return get_disk_usage(study_path) if study_path.exists() else 0

antarest/study/storage/utils.py

@@ -230,7 +230,12 @@ def assert_permission(
         permission_type: level of permission
         raising: raise error if permission not matched
 
-    Returns: true if permission match, false if not raising.
+    Returns:
+        `True` if the user has the required permissions, `False` otherwise.
+
+    Raises:
+        `UserHasNotPermissionError`: If the raising parameter is set to `True`
+            and the user does not have the required permissions.
     """
     studies = [study] if study else []
     return assert_permission_on_studies(user, studies, permission_type, raising=raising)

tests/integration/studies_blueprint/test_disk_usage.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 from starlette.testclient import TestClient
 
 from antarest.core.tasks.model import TaskDTO, TaskStatus
@@ -9,6 +11,7 @@ def test_disk_usage_endpoint(
         client: TestClient,
         user_access_token: str,
         study_id: str,
+        tmp_path: Path,
     ) -> None:
         """
         Verify the functionality of the disk usage endpoint:
@@ -63,11 +66,12 @@ def test_disk_usage_endpoint(
         res.raise_for_status()
         task_id = res.json()
 
-        # wait for task completion
-        res = client.get(f"/v1/tasks/{task_id}?wait_for_completion=true", headers=user_headers)
+        # Wait for task completion
+        res = client.get(f"/v1/tasks/{task_id}", headers=user_headers, params={"wait_for_completion": True})
         assert res.status_code == 200
         task_result = TaskDTO.parse_obj(res.json())
         assert task_result.status == TaskStatus.COMPLETED
+        assert task_result.result is not None
         assert task_result.result.success
 
         # Ensure a successful response is received and the disk usage is not zero
@@ -78,3 +82,19 @@ def test_disk_usage_endpoint(
         assert res.status_code == 200, res.json()
         disk_usage = res.json()  # currently: 6.38 Mio on Ubuntu.
         assert 6 * 1024 * 1024 < disk_usage < 7 * 1024 * 1024
+
+        # Create dummy outputs for the variant
+        outputs_dir = tmp_path / "internal_workspace" / variant_id / "output"
+        for output in ["20240208-1200eco", "20240208-1300eco"]:
+            output_dir = outputs_dir / output
+            output_dir.mkdir(parents=True)
+            (output_dir / "dummy.txt").write_bytes(b"dummy content")
+
+        # Calculate the new disk usage
+        res = client.get(
+            f"/v1/studies/{variant_id}/disk-usage",
+            headers=user_headers,
+        )
+        assert res.status_code == 200, res.json()
+        new_disk_usage = res.json()
+        assert new_disk_usage == disk_usage + 2 * len(b"dummy content")