diff --git a/antarest/study/service.py b/antarest/study/service.py index b80b0b4438..ae86fe62ae 100644 --- a/antarest/study/service.py +++ b/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 diff --git a/antarest/study/storage/utils.py b/antarest/study/storage/utils.py index b23d10333a..a0fc7a02fe 100644 --- a/antarest/study/storage/utils.py +++ b/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) diff --git a/docs/assets/media/img/userguide_studyoverview.png b/docs/assets/media/img/userguide_studyoverview.png index a6de2e7ed0..f7516c8c9b 100644 Binary files a/docs/assets/media/img/userguide_studyoverview.png and b/docs/assets/media/img/userguide_studyoverview.png differ diff --git a/tests/integration/studies_blueprint/test_disk_usage.py b/tests/integration/studies_blueprint/test_disk_usage.py index 95e3afbad8..1685acad87 100644 --- a/tests/integration/studies_blueprint/test_disk_usage.py +++ b/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")