From 5e1f1cfccdc4443fa61b3d19bdb0888446ca966e Mon Sep 17 00:00:00 2001 From: Gyubong Date: Wed, 28 Jan 2026 11:02:32 +0900 Subject: [PATCH 1/3] wip --- src/ai/backend/manager/errors/vfolder.py | 67 +++++++++++++++ .../services/artifact_revision/service.py | 82 ++++++++++++++++++- 2 files changed, 148 insertions(+), 1 deletion(-) create mode 100644 src/ai/backend/manager/errors/vfolder.py diff --git a/src/ai/backend/manager/errors/vfolder.py b/src/ai/backend/manager/errors/vfolder.py new file mode 100644 index 00000000000..6f30f8c31ea --- /dev/null +++ b/src/ai/backend/manager/errors/vfolder.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from aiohttp import web + +from ai.backend.common.exception import ( + BackendAIError, + ErrorCode, + ErrorDetail, + ErrorDomain, + ErrorOperation, +) + +if TYPE_CHECKING: + from ai.backend.common.types import QuotaScopeID, VFolderID + + +class VFolderQuotaExceededError(BackendAIError, web.HTTPBadRequest): + """ + Raised when a VFolder quota scope limit would be exceeded by an operation. + + This error is raised during pre-validation before operations like artifact import, + preventing partial downloads and wasted resources. + """ + + error_type = "https://api.backend.ai/probs/vfolder-quota-exceeded" + error_title = "VFolder Quota Scope Limit Would Be Exceeded" + + def __init__( + self, + vfolder_id: VFolderID, + quota_scope_id: QuotaScopeID, + current_size: int, + max_size: int, + requested_size: int, + ) -> None: + self.vfolder_id = vfolder_id + self.quota_scope_id = quota_scope_id + self.current_size = current_size + self.max_size = max_size + self.requested_size = requested_size + available_bytes = max_size - current_size + + message = ( + f"VFolder quota scope limit would be exceeded. " + f"Current: {current_size} bytes, Max: {max_size} bytes, " + f"Requested: {requested_size} bytes, Available: {available_bytes} bytes" + ) + super().__init__(message) + + def error_code(self) -> ErrorCode: + return ErrorCode( + domain=ErrorDomain.QUOTA_SCOPE, + operation=ErrorOperation.CHECK_LIMIT, + error_detail=ErrorDetail.BAD_REQUEST, + ) + + def error_data(self) -> dict: + return { + "vfolder_id": str(self.vfolder_id), + "quota_scope_id": str(self.quota_scope_id), + "current_size_bytes": self.current_size, + "max_size_bytes": self.max_size, + "requested_size_bytes": self.requested_size, + "available_bytes": self.max_size - self.current_size, + } diff --git a/src/ai/backend/manager/services/artifact_revision/service.py b/src/ai/backend/manager/services/artifact_revision/service.py index 054e0d3ff13..ef7afcf3cb7 100644 --- a/src/ai/backend/manager/services/artifact_revision/service.py +++ b/src/ai/backend/manager/services/artifact_revision/service.py @@ -29,7 +29,7 @@ HuggingFaceImportModelsReq, ReservoirImportModelsReq, ) -from ai.backend.common.types import VFolderID +from ai.backend.common.types import QuotaScopeID, QuotaScopeType, VFolderID from ai.backend.logging import BraceStyleAdapter from ai.backend.manager.clients.artifact_registry.reservoir_client import ReservoirRegistryClient from ai.backend.manager.clients.storage_proxy.session_manager import StorageSessionManager @@ -44,6 +44,7 @@ ArtifactType, ) from ai.backend.manager.data.artifact_registries.types import ArtifactRegistryData +from ai.backend.manager.data.vfolder.types import VFolderData from ai.backend.manager.dto.request import ( DelegateImportArtifactsReq, ImportArtifactsOptions, @@ -59,6 +60,7 @@ ) from ai.backend.manager.errors.common import ServerMisconfiguredError from ai.backend.manager.errors.storage import UnsupportedStorageTypeError +from ai.backend.manager.errors.vfolder import VFolderQuotaExceededError from ai.backend.manager.repositories.artifact.repository import ArtifactRepository from ai.backend.manager.repositories.artifact_registry.repository import ArtifactRegistryRepository from ai.backend.manager.repositories.huggingface_registry.repository import HuggingFaceRepository @@ -422,6 +424,10 @@ async def import_revision( vfolder_id = VFolderID(vfolder_data.quota_scope_id, vfolder_data.id) _, volume_name = self._storage_manager.get_proxy_and_volume(vfolder_data.host) + # Pre-validate VFolder quota before starting import + if revision_data.size is not None: + await self._check_vfolder_quota(vfolder_data, revision_data.size) + task_id: UUID match artifact.registry_type: case ArtifactRegistryType.HUGGINGFACE: @@ -812,3 +818,77 @@ def _is_latest_commit_hash( artifact_revision.status == ArtifactStatus.AVAILABLE and artifact_revision.digest == latest_commit_hash ) + + async def _check_vfolder_quota( + self, + vfolder_data: VFolderData, + additional_size: int, + ) -> None: + """ + Check the quota scope limit for a VFolder before importing artifacts. + + This pre-validation prevents failed imports mid-way when storage proxy + detects quota exceeded, avoiding partial downloads. + + :param vfolder_data: VFolder data containing quota_scope_id and host + :param additional_size: Size of the artifact to be imported (in bytes) + :raises VFolderQuotaExceededError: If the quota limit would be exceeded + """ + quota_scope_id = vfolder_data.quota_scope_id + if quota_scope_id is None: + # No quota scope configured for this VFolder + return + + _, volume_name = self._storage_manager.get_proxy_and_volume(vfolder_data.host) + + # 1. Get current quota scope usage from storage proxy + storage_client = self._storage_manager.get_manager_facing_client(vfolder_data.host) + usage_response = await storage_client.get_quota_scope(volume_name, str(quota_scope_id)) + used_bytes = int(usage_response.get("used_bytes", 0)) + + # 2. Get limit from resource policy + max_size = await self._get_quota_scope_limit(quota_scope_id, vfolder_data.domain_name) + + # -1 or 0 means unlimited + if max_size <= 0: + return + + # 3. Check if limit would be exceeded + if used_bytes + additional_size > max_size: + vfolder_id = VFolderID(quota_scope_id, vfolder_data.id) + raise VFolderQuotaExceededError( + vfolder_id=vfolder_id, + quota_scope_id=quota_scope_id, + current_size=used_bytes, + max_size=max_size, + requested_size=additional_size, + ) + + async def _get_quota_scope_limit(self, quota_scope_id: QuotaScopeID, domain_name: str) -> int: + """ + Get the quota scope limit from the appropriate resource policy. + + :param quota_scope_id: The quota scope ID to get the limit for + :param domain_name: The domain name for project scope queries + :return: The max_quota_scope_size in bytes, or -1 if unlimited + """ + match quota_scope_id.scope_type: + case QuotaScopeType.USER: + user_result = await self._vfolder_repository.get_user_resource_info( + quota_scope_id.scope_id + ) + if user_result is None: + # User not found, treat as unlimited + return -1 + _, max_quota_scope_size, _ = user_result + return max_quota_scope_size + + case QuotaScopeType.PROJECT: + group_result = await self._vfolder_repository.get_group_resource_info( + quota_scope_id.scope_id, domain_name + ) + if group_result is None: + # Group not found, treat as unlimited + return -1 + _, _, max_quota_scope_size, _ = group_result + return max_quota_scope_size From 8c53223fe4b1f6389974bf9e253aafe9fb7d172c Mon Sep 17 00:00:00 2001 From: Gyubong Date: Fri, 30 Jan 2026 17:22:54 +0900 Subject: [PATCH 2/3] test: Add --- changes/8390.feature.md | 1 + .../test_artifact_revision_service.py | 491 ++++++++++++++++++ 2 files changed, 492 insertions(+) create mode 100644 changes/8390.feature.md diff --git a/changes/8390.feature.md b/changes/8390.feature.md new file mode 100644 index 00000000000..2d6c857316d --- /dev/null +++ b/changes/8390.feature.md @@ -0,0 +1 @@ +Implement artifact storage usage tracking and quota enforcement diff --git a/tests/unit/manager/services/artifact_revision/test_artifact_revision_service.py b/tests/unit/manager/services/artifact_revision/test_artifact_revision_service.py index 665716bdb37..a76418cb701 100644 --- a/tests/unit/manager/services/artifact_revision/test_artifact_revision_service.py +++ b/tests/unit/manager/services/artifact_revision/test_artifact_revision_service.py @@ -6,6 +6,7 @@ from __future__ import annotations +import dataclasses import uuid from datetime import UTC, datetime from unittest.mock import AsyncMock, MagicMock @@ -13,6 +14,7 @@ import pytest from ai.backend.common.data.artifact.types import ArtifactRegistryType +from ai.backend.common.types import QuotaScopeID, QuotaScopeType, VFolderID, VFolderUsageMode from ai.backend.manager.data.artifact.types import ( ArtifactAvailability, ArtifactData, @@ -23,6 +25,13 @@ ArtifactStatus, ArtifactType, ) +from ai.backend.manager.data.vfolder.types import ( + VFolderData, + VFolderMountPermission, + VFolderOperationStatus, + VFolderOwnershipType, +) +from ai.backend.manager.errors.vfolder import VFolderQuotaExceededError from ai.backend.manager.repositories.artifact.repository import ArtifactRepository from ai.backend.manager.repositories.artifact_registry.repository import ArtifactRegistryRepository from ai.backend.manager.repositories.base import BatchQuerier, OffsetPagination @@ -528,3 +537,485 @@ async def test_upsert_artifacts_with_revisions( assert result.result[0].revisions == [sample_artifact_revision] mock_artifact_repository.upsert_artifacts.assert_called_once() mock_artifact_repository.upsert_artifact_revisions.assert_called_once() + + +class TestVFolderQuotaCheck: + """Tests for VFolder quota validation in ArtifactRevisionService.""" + + @pytest.fixture + def mock_artifact_repository(self) -> MagicMock: + return MagicMock(spec=ArtifactRepository) + + @pytest.fixture + def mock_artifact_registry_repository(self) -> MagicMock: + return MagicMock(spec=ArtifactRegistryRepository) + + @pytest.fixture + def mock_object_storage_repository(self) -> MagicMock: + return MagicMock(spec=ObjectStorageRepository) + + @pytest.fixture + def mock_vfs_storage_repository(self) -> MagicMock: + return MagicMock(spec=VFSStorageRepository) + + @pytest.fixture + def mock_storage_namespace_repository(self) -> MagicMock: + return MagicMock(spec=StorageNamespaceRepository) + + @pytest.fixture + def mock_huggingface_repository(self) -> MagicMock: + return MagicMock(spec=HuggingFaceRepository) + + @pytest.fixture + def mock_reservoir_repository(self) -> MagicMock: + return MagicMock(spec=ReservoirRegistryRepository) + + @pytest.fixture + def mock_vfolder_repository(self) -> MagicMock: + return MagicMock(spec=VfolderRepository) + + @pytest.fixture + def mock_storage_manager(self) -> MagicMock: + return MagicMock() + + @pytest.fixture + def mock_config_provider(self) -> MagicMock: + return MagicMock() + + @pytest.fixture + def mock_valkey_artifact_client(self) -> MagicMock: + return MagicMock() + + @pytest.fixture + def mock_background_task_manager(self) -> MagicMock: + return MagicMock() + + @pytest.fixture + def artifact_revision_service( + self, + mock_artifact_repository: MagicMock, + mock_artifact_registry_repository: MagicMock, + mock_object_storage_repository: MagicMock, + mock_vfs_storage_repository: MagicMock, + mock_storage_namespace_repository: MagicMock, + mock_huggingface_repository: MagicMock, + mock_reservoir_repository: MagicMock, + mock_vfolder_repository: MagicMock, + mock_storage_manager: MagicMock, + mock_config_provider: MagicMock, + mock_valkey_artifact_client: MagicMock, + mock_background_task_manager: MagicMock, + ) -> ArtifactRevisionService: + return ArtifactRevisionService( + artifact_repository=mock_artifact_repository, + artifact_registry_repository=mock_artifact_registry_repository, + object_storage_repository=mock_object_storage_repository, + vfs_storage_repository=mock_vfs_storage_repository, + storage_namespace_repository=mock_storage_namespace_repository, + huggingface_registry_repository=mock_huggingface_repository, + reservoir_registry_repository=mock_reservoir_repository, + vfolder_repository=mock_vfolder_repository, + storage_manager=mock_storage_manager, + config_provider=mock_config_provider, + valkey_artifact_client=mock_valkey_artifact_client, + background_task_manager=mock_background_task_manager, + ) + + @pytest.fixture + def sample_vfolder_data(self) -> VFolderData: + """Create sample VFolderData with quota scope.""" + return VFolderData( + id=uuid.uuid4(), + name="test-vfolder", + host="local:volume1", + domain_name="default", + quota_scope_id=QuotaScopeID(QuotaScopeType.USER, uuid.uuid4()), + usage_mode=VFolderUsageMode.MODEL, + permission=VFolderMountPermission.RW_DELETE, + max_files=1000, + max_size=1024 * 1024 * 1024, + num_files=10, + cur_size=100 * 1024 * 1024, + created_at=datetime.now(UTC), + last_used=None, + creator="testuser@example.com", + unmanaged_path=None, + ownership_type=VFolderOwnershipType.USER, + user=uuid.uuid4(), + group=None, + cloneable=False, + status=VFolderOperationStatus.READY, + ) + + async def test_check_vfolder_quota_no_quota_scope( + self, + artifact_revision_service: ArtifactRevisionService, + sample_vfolder_data: VFolderData, + ) -> None: + """When quota_scope_id is None, quota check should be skipped.""" + vfolder_data = dataclasses.replace(sample_vfolder_data, quota_scope_id=None) + + # Should not raise any error + await artifact_revision_service._check_vfolder_quota(vfolder_data, 500 * 1024 * 1024) + + async def test_check_vfolder_quota_unlimited_negative_one( + self, + artifact_revision_service: ArtifactRevisionService, + mock_storage_manager: MagicMock, + mock_vfolder_repository: MagicMock, + sample_vfolder_data: VFolderData, + ) -> None: + """When max quota is -1 (unlimited), should not raise error.""" + mock_storage_manager.get_proxy_and_volume.return_value = ("proxy", "volume1") + mock_storage_client = AsyncMock() + mock_storage_client.get_quota_scope = AsyncMock( + return_value={"used_bytes": 500 * 1024 * 1024} + ) + mock_storage_manager.get_manager_facing_client.return_value = mock_storage_client + + mock_vfolder_repository.get_user_resource_info = AsyncMock( + return_value=("user_uuid", -1, "policy") + ) + + # Should not raise any error even with large additional size + await artifact_revision_service._check_vfolder_quota( + sample_vfolder_data, 10 * 1024 * 1024 * 1024 + ) + + async def test_check_vfolder_quota_unlimited_zero( + self, + artifact_revision_service: ArtifactRevisionService, + mock_storage_manager: MagicMock, + mock_vfolder_repository: MagicMock, + sample_vfolder_data: VFolderData, + ) -> None: + """When max quota is 0 (unlimited), should not raise error.""" + mock_storage_manager.get_proxy_and_volume.return_value = ("proxy", "volume1") + mock_storage_client = AsyncMock() + mock_storage_client.get_quota_scope = AsyncMock( + return_value={"used_bytes": 500 * 1024 * 1024} + ) + mock_storage_manager.get_manager_facing_client.return_value = mock_storage_client + + mock_vfolder_repository.get_user_resource_info = AsyncMock( + return_value=("user_uuid", 0, "policy") + ) + + await artifact_revision_service._check_vfolder_quota( + sample_vfolder_data, 10 * 1024 * 1024 * 1024 + ) + + async def test_check_vfolder_quota_within_limit( + self, + artifact_revision_service: ArtifactRevisionService, + mock_storage_manager: MagicMock, + mock_vfolder_repository: MagicMock, + sample_vfolder_data: VFolderData, + ) -> None: + """When within quota limit, should not raise error.""" + mock_storage_manager.get_proxy_and_volume.return_value = ("proxy", "volume1") + mock_storage_client = AsyncMock() + mock_storage_client.get_quota_scope = AsyncMock( + return_value={"used_bytes": 100 * 1024 * 1024} + ) + mock_storage_manager.get_manager_facing_client.return_value = mock_storage_client + + mock_vfolder_repository.get_user_resource_info = AsyncMock( + return_value=("user_uuid", 1024 * 1024 * 1024, "policy") + ) + + # 100MB (used) + 500MB (additional) < 1GB (max) → OK + await artifact_revision_service._check_vfolder_quota(sample_vfolder_data, 500 * 1024 * 1024) + + async def test_check_vfolder_quota_exactly_at_limit( + self, + artifact_revision_service: ArtifactRevisionService, + mock_storage_manager: MagicMock, + mock_vfolder_repository: MagicMock, + sample_vfolder_data: VFolderData, + ) -> None: + """When exactly at quota limit, should not raise error.""" + mock_storage_manager.get_proxy_and_volume.return_value = ("proxy", "volume1") + mock_storage_client = AsyncMock() + mock_storage_client.get_quota_scope = AsyncMock( + return_value={"used_bytes": 500 * 1024 * 1024} + ) + mock_storage_manager.get_manager_facing_client.return_value = mock_storage_client + + mock_vfolder_repository.get_user_resource_info = AsyncMock( + return_value=("user_uuid", 1024 * 1024 * 1024, "policy") + ) + + # 500MB (used) + 524MB (additional) = 1024MB (max) → OK (exactly at limit) + await artifact_revision_service._check_vfolder_quota(sample_vfolder_data, 524 * 1024 * 1024) + + async def test_check_vfolder_quota_exceeded( + self, + artifact_revision_service: ArtifactRevisionService, + mock_storage_manager: MagicMock, + mock_vfolder_repository: MagicMock, + sample_vfolder_data: VFolderData, + ) -> None: + """When quota would be exceeded, should raise VFolderQuotaExceededError.""" + mock_storage_manager.get_proxy_and_volume.return_value = ("proxy", "volume1") + mock_storage_client = AsyncMock() + mock_storage_client.get_quota_scope = AsyncMock( + return_value={"used_bytes": 900 * 1024 * 1024} + ) + mock_storage_manager.get_manager_facing_client.return_value = mock_storage_client + + mock_vfolder_repository.get_user_resource_info = AsyncMock( + return_value=("user_uuid", 1024 * 1024 * 1024, "policy") + ) + + # 900MB (used) + 200MB (additional) > 1GB (max) → Error + with pytest.raises(VFolderQuotaExceededError) as exc_info: + await artifact_revision_service._check_vfolder_quota( + sample_vfolder_data, 200 * 1024 * 1024 + ) + + error = exc_info.value + assert error.current_size == 900 * 1024 * 1024 + assert error.max_size == 1024 * 1024 * 1024 + assert error.requested_size == 200 * 1024 * 1024 + + async def test_check_vfolder_quota_project_scope( + self, + artifact_revision_service: ArtifactRevisionService, + mock_storage_manager: MagicMock, + mock_vfolder_repository: MagicMock, + sample_vfolder_data: VFolderData, + ) -> None: + """Test quota check with PROJECT scope type.""" + project_id = uuid.uuid4() + vfolder_data = dataclasses.replace( + sample_vfolder_data, + quota_scope_id=QuotaScopeID(QuotaScopeType.PROJECT, project_id), + ) + + mock_storage_manager.get_proxy_and_volume.return_value = ("proxy", "volume1") + mock_storage_client = AsyncMock() + mock_storage_client.get_quota_scope = AsyncMock( + return_value={"used_bytes": 100 * 1024 * 1024} + ) + mock_storage_manager.get_manager_facing_client.return_value = mock_storage_client + + mock_vfolder_repository.get_group_resource_info = AsyncMock( + return_value=("group_name", "domain", 2 * 1024 * 1024 * 1024, "policy") + ) + + # Should not raise error + await artifact_revision_service._check_vfolder_quota(vfolder_data, 500 * 1024 * 1024) + + mock_vfolder_repository.get_group_resource_info.assert_called_once_with( + project_id, "default" + ) + + +class TestGetQuotaScopeLimit: + """Tests for _get_quota_scope_limit method.""" + + @pytest.fixture + def mock_artifact_repository(self) -> MagicMock: + return MagicMock(spec=ArtifactRepository) + + @pytest.fixture + def mock_artifact_registry_repository(self) -> MagicMock: + return MagicMock(spec=ArtifactRegistryRepository) + + @pytest.fixture + def mock_object_storage_repository(self) -> MagicMock: + return MagicMock(spec=ObjectStorageRepository) + + @pytest.fixture + def mock_vfs_storage_repository(self) -> MagicMock: + return MagicMock(spec=VFSStorageRepository) + + @pytest.fixture + def mock_storage_namespace_repository(self) -> MagicMock: + return MagicMock(spec=StorageNamespaceRepository) + + @pytest.fixture + def mock_huggingface_repository(self) -> MagicMock: + return MagicMock(spec=HuggingFaceRepository) + + @pytest.fixture + def mock_reservoir_repository(self) -> MagicMock: + return MagicMock(spec=ReservoirRegistryRepository) + + @pytest.fixture + def mock_vfolder_repository(self) -> MagicMock: + return MagicMock(spec=VfolderRepository) + + @pytest.fixture + def mock_storage_manager(self) -> MagicMock: + return MagicMock() + + @pytest.fixture + def mock_config_provider(self) -> MagicMock: + return MagicMock() + + @pytest.fixture + def mock_valkey_artifact_client(self) -> MagicMock: + return MagicMock() + + @pytest.fixture + def mock_background_task_manager(self) -> MagicMock: + return MagicMock() + + @pytest.fixture + def artifact_revision_service( + self, + mock_artifact_repository: MagicMock, + mock_artifact_registry_repository: MagicMock, + mock_object_storage_repository: MagicMock, + mock_vfs_storage_repository: MagicMock, + mock_storage_namespace_repository: MagicMock, + mock_huggingface_repository: MagicMock, + mock_reservoir_repository: MagicMock, + mock_vfolder_repository: MagicMock, + mock_storage_manager: MagicMock, + mock_config_provider: MagicMock, + mock_valkey_artifact_client: MagicMock, + mock_background_task_manager: MagicMock, + ) -> ArtifactRevisionService: + return ArtifactRevisionService( + artifact_repository=mock_artifact_repository, + artifact_registry_repository=mock_artifact_registry_repository, + object_storage_repository=mock_object_storage_repository, + vfs_storage_repository=mock_vfs_storage_repository, + storage_namespace_repository=mock_storage_namespace_repository, + huggingface_registry_repository=mock_huggingface_repository, + reservoir_registry_repository=mock_reservoir_repository, + vfolder_repository=mock_vfolder_repository, + storage_manager=mock_storage_manager, + config_provider=mock_config_provider, + valkey_artifact_client=mock_valkey_artifact_client, + background_task_manager=mock_background_task_manager, + ) + + async def test_get_quota_scope_limit_user_type( + self, + artifact_revision_service: ArtifactRevisionService, + mock_vfolder_repository: MagicMock, + ) -> None: + """Test _get_quota_scope_limit with USER scope type.""" + user_id = uuid.uuid4() + quota_scope_id = QuotaScopeID(QuotaScopeType.USER, user_id) + + mock_vfolder_repository.get_user_resource_info = AsyncMock( + return_value=("user_uuid", 1024 * 1024 * 1024, "policy") + ) + + result = await artifact_revision_service._get_quota_scope_limit(quota_scope_id, "default") + + assert result == 1024 * 1024 * 1024 + mock_vfolder_repository.get_user_resource_info.assert_called_once_with(user_id) + + async def test_get_quota_scope_limit_project_type( + self, + artifact_revision_service: ArtifactRevisionService, + mock_vfolder_repository: MagicMock, + ) -> None: + """Test _get_quota_scope_limit with PROJECT scope type.""" + project_id = uuid.uuid4() + quota_scope_id = QuotaScopeID(QuotaScopeType.PROJECT, project_id) + + mock_vfolder_repository.get_group_resource_info = AsyncMock( + return_value=("group_name", "domain", 2 * 1024 * 1024 * 1024, "policy") + ) + + result = await artifact_revision_service._get_quota_scope_limit(quota_scope_id, "default") + + assert result == 2 * 1024 * 1024 * 1024 + mock_vfolder_repository.get_group_resource_info.assert_called_once_with( + project_id, "default" + ) + + async def test_get_quota_scope_limit_user_not_found( + self, + artifact_revision_service: ArtifactRevisionService, + mock_vfolder_repository: MagicMock, + ) -> None: + """When user is not found, should return -1 (unlimited).""" + user_id = uuid.uuid4() + quota_scope_id = QuotaScopeID(QuotaScopeType.USER, user_id) + + mock_vfolder_repository.get_user_resource_info = AsyncMock(return_value=None) + + result = await artifact_revision_service._get_quota_scope_limit(quota_scope_id, "default") + + assert result == -1 + + async def test_get_quota_scope_limit_project_not_found( + self, + artifact_revision_service: ArtifactRevisionService, + mock_vfolder_repository: MagicMock, + ) -> None: + """When project is not found, should return -1 (unlimited).""" + project_id = uuid.uuid4() + quota_scope_id = QuotaScopeID(QuotaScopeType.PROJECT, project_id) + + mock_vfolder_repository.get_group_resource_info = AsyncMock(return_value=None) + + result = await artifact_revision_service._get_quota_scope_limit(quota_scope_id, "default") + + assert result == -1 + + +class TestVFolderQuotaExceededError: + """Tests for VFolderQuotaExceededError.""" + + def test_error_message_format(self) -> None: + """Test error message contains all relevant information.""" + quota_scope_id = QuotaScopeID(QuotaScopeType.USER, uuid.uuid4()) + vfolder_id = VFolderID(quota_scope_id, uuid.uuid4()) + + error = VFolderQuotaExceededError( + vfolder_id=vfolder_id, + quota_scope_id=quota_scope_id, + current_size=900, + max_size=1000, + requested_size=200, + ) + + error_message = str(error) + assert "900" in error_message + assert "1000" in error_message + assert "200" in error_message + assert "100" in error_message # available bytes + + def test_error_data_structure(self) -> None: + """Test error_data returns correct structure.""" + quota_scope_id = QuotaScopeID(QuotaScopeType.USER, uuid.uuid4()) + vfolder_id = VFolderID(quota_scope_id, uuid.uuid4()) + + error = VFolderQuotaExceededError( + vfolder_id=vfolder_id, + quota_scope_id=quota_scope_id, + current_size=900, + max_size=1000, + requested_size=200, + ) + + data = error.error_data() + assert data["current_size_bytes"] == 900 + assert data["max_size_bytes"] == 1000 + assert data["requested_size_bytes"] == 200 + assert data["available_bytes"] == 100 + + def test_error_code(self) -> None: + """Test error_code returns correct ErrorCode.""" + quota_scope_id = QuotaScopeID(QuotaScopeType.USER, uuid.uuid4()) + vfolder_id = VFolderID(quota_scope_id, uuid.uuid4()) + + error = VFolderQuotaExceededError( + vfolder_id=vfolder_id, + quota_scope_id=quota_scope_id, + current_size=900, + max_size=1000, + requested_size=200, + ) + + error_code = error.error_code() + assert error_code is not None From f8f10dd20e6344c711c875524183178895d0418f Mon Sep 17 00:00:00 2001 From: Gyubong Date: Wed, 4 Feb 2026 11:16:03 +0900 Subject: [PATCH 3/3] fix: Merge --- src/ai/backend/manager/errors/vfolder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ai/backend/manager/errors/vfolder.py b/src/ai/backend/manager/errors/vfolder.py index 6f30f8c31ea..de1a21b65a1 100644 --- a/src/ai/backend/manager/errors/vfolder.py +++ b/src/ai/backend/manager/errors/vfolder.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from aiohttp import web @@ -56,7 +56,7 @@ def error_code(self) -> ErrorCode: error_detail=ErrorDetail.BAD_REQUEST, ) - def error_data(self) -> dict: + def error_data(self) -> dict[str, Any]: return { "vfolder_id": str(self.vfolder_id), "quota_scope_id": str(self.quota_scope_id),