From 00b3876093db1d6b150bbbaafef37bffe2f86640 Mon Sep 17 00:00:00 2001
From: MartinBelthle <102529366+MartinBelthle@users.noreply.github.com>
Date: Wed, 31 Jan 2024 17:43:43 +0100
Subject: [PATCH] fix(disk-usage): fix bug on variants (#1915)

---
 antarest/study/service.py                     |  5 +-
 .../studies_blueprint/test_disk_usage.py      | 59 ++++++++++++++++++-
 2 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/antarest/study/service.py b/antarest/study/service.py
index 98d36dee9b..79cdb930fe 100644
--- a/antarest/study/service.py
+++ b/antarest/study/service.py
@@ -2373,5 +2373,6 @@ def get_disk_usage(self, uuid: str, params: RequestParameters) -> int:
         """
         study = self.get_study(uuid=uuid)
         assert_permission(params.user, study, StudyPermissionType.READ)
-        path = str(self.storage_service.get_storage(study).get_study_path(study))
-        return get_disk_usage(path=path)
+        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
diff --git a/tests/integration/studies_blueprint/test_disk_usage.py b/tests/integration/studies_blueprint/test_disk_usage.py
index 89d3b70e74..95e3afbad8 100644
--- a/tests/integration/studies_blueprint/test_disk_usage.py
+++ b/tests/integration/studies_blueprint/test_disk_usage.py
@@ -1,5 +1,7 @@
 from starlette.testclient import TestClient
 
+from antarest.core.tasks.model import TaskDTO, TaskStatus
+
 
 class TestDiskUsage:
     def test_disk_usage_endpoint(
@@ -14,10 +16,65 @@ def test_disk_usage_endpoint(
         - Ensure a successful response is received.
         - Confirm that the JSON response is an integer which represent a (big enough) directory size.
         """
+        disk_usage: int
+
+        user_headers = {"Authorization": f"Bearer {user_access_token}"}
         res = client.get(
             f"/v1/studies/{study_id}/disk-usage",
-            headers={"Authorization": f"Bearer {user_access_token}"},
+            headers=user_headers,
         )
         assert res.status_code == 200, res.json()
         disk_usage = res.json()  # currently: 7.47 Mio on Ubuntu
         assert 7 * 1024 * 1024 < disk_usage < 8 * 1024 * 1024
+
+        # Copy the study in managed workspace in order to create a variant
+        res = client.post(
+            f"/v1/studies/{study_id}/copy",
+            headers=user_headers,
+            params={"dest": "somewhere", "use_task": "false"},
+        )
+        res.raise_for_status()
+        parent_id: str = res.json()
+
+        # Create variant of the copied study
+        res = client.post(
+            f"/v1/studies/{parent_id}/variants",
+            headers=user_headers,
+            params={"name": "Variant Test"},
+        )
+        res.raise_for_status()
+        variant_id: str = res.json()
+
+        # Ensure a successful response is received even if the variant has no snapshot
+        res = client.get(
+            f"/v1/studies/{variant_id}/disk-usage",
+            headers=user_headers,
+        )
+        assert res.status_code == 200, res.json()
+        disk_usage = res.json()
+        assert disk_usage == 0
+
+        # Generate a snapshot for the variant
+        res = client.put(
+            f"/v1/studies/{variant_id}/generate",
+            headers={"Authorization": f"Bearer {user_access_token}"},
+            params={"denormalize": True, "from_scratch": True},
+        )
+        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)
+        assert res.status_code == 200
+        task_result = TaskDTO.parse_obj(res.json())
+        assert task_result.status == TaskStatus.COMPLETED
+        assert task_result.result.success
+
+        # Ensure a successful response is received and the disk usage is not zero
+        res = client.get(
+            f"/v1/studies/{variant_id}/disk-usage",
+            headers=user_headers,
+        )
+        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