diff --git a/packages/domain/quota.py b/packages/domain/quota.py new file mode 100644 index 000000000..aa20811f8 --- /dev/null +++ b/packages/domain/quota.py @@ -0,0 +1,178 @@ +""" +配额检查服务 +用于检查工作空间是否超出配额限制 +""" +from typing import Optional + + +class QuotaChecker: + """配额检查器""" + + def __init__( + self, + workspace_repository, + project_repository, + ): + self.workspace_repository = workspace_repository + self.project_repository = project_repository + + def check_can_create_project( + self, + workspace_id: str, + ) -> tuple[bool, Optional[str]]: + """ + 检查是否可以创建项目 + + Args: + workspace_id: 工作空间 ID + + Returns: + (是否可以, 错误信息) + """ + workspace = self.workspace_repository.find_by_id(workspace_id) + if not workspace: + return False, "Workspace not found" + + # 获取当前项目数量 + current_count = self.project_repository.count_by_workspace(workspace_id) + + # 检查是否超出配额(999999 表示无限) + if workspace.max_projects != 999999 and current_count >= workspace.max_projects: + return False, f"Project limit reached ({workspace.max_projects}). Upgrade your plan to create more projects." + + return True, None + + def check_storage_available( + self, + workspace_id: str, + additional_gb: float, + ) -> tuple[bool, Optional[str]]: + """ + 检查存储空间是否足够 + + Args: + workspace_id: 工作空间 ID + additional_gb: 需要的额外存储空间(GB) + + Returns: + (是否可以, 错误信息) + """ + workspace = self.workspace_repository.find_by_id(workspace_id) + if not workspace: + return False, "Workspace not found" + + # 检查存储空间 + new_usage = workspace.used_storage_gb + additional_gb + + if new_usage > workspace.max_storage_gb: + remaining = workspace.max_storage_gb - workspace.used_storage_gb + return False, f"Storage limit exceeded. Available: {remaining:.2f}GB, Required: {additional_gb:.2f}GB. Upgrade your plan for more storage." + + return True, None + + def get_quota_status(self, workspace_id: str) -> dict: + """ + 获取配额使用状态 + + Args: + workspace_id: 工作空间 ID + + Returns: + 配额状态信息 + """ + workspace = self.workspace_repository.find_by_id(workspace_id) + if not workspace: + return None + + # 获取项目数量 + project_count = self.project_repository.count_by_workspace(workspace_id) + + # 计算使用率 + project_usage_percent = ( + (project_count / workspace.max_projects * 100) + if workspace.max_projects != 999999 + else 0 # 无限制 + ) + + storage_usage_percent = ( + (workspace.used_storage_gb / workspace.max_storage_gb * 100) + if workspace.max_storage_gb > 0 + else 0 + ) + + return { + "workspace_id": workspace.id, + "subscription_plan": workspace.subscription_plan, + "projects": { + "used": project_count, + "limit": workspace.max_projects, + "unlimited": workspace.max_projects == 999999, + "usage_percent": project_usage_percent, + }, + "storage": { + "used_gb": workspace.used_storage_gb, + "limit_gb": workspace.max_storage_gb, + "remaining_gb": workspace.max_storage_gb - workspace.used_storage_gb, + "usage_percent": storage_usage_percent, + }, + } + + def update_storage_usage( + self, + workspace_id: str, + delta_gb: float, + ) -> tuple[bool, Optional[str]]: + """ + 更新存储使用量 + + Args: + workspace_id: 工作空间 ID + delta_gb: 变化量(正数为增加,负数为减少) + + Returns: + (是否成功, 错误信息) + """ + workspace = self.workspace_repository.find_by_id(workspace_id) + if not workspace: + return False, "Workspace not found" + + # 更新使用量 + new_usage = workspace.used_storage_gb + delta_gb + + # 不能为负数 + if new_usage < 0: + new_usage = 0 + + workspace.used_storage_gb = new_usage + self.workspace_repository.save(workspace) + + return True, None + + +class QuotaWarningLevel: + """配额警告级别""" + + NORMAL = "normal" # <80% + WARNING = "warning" # 80-90% + CRITICAL = "critical" # 90-100% + EXCEEDED = "exceeded" # >100% + + +def get_warning_level(usage_percent: float) -> str: + """ + 根据使用率获取警告级别 + + Args: + usage_percent: 使用率(0-100) + + Returns: + 警告级别 + """ + if usage_percent >= 100: + return QuotaWarningLevel.EXCEEDED + elif usage_percent >= 90: + return QuotaWarningLevel.CRITICAL + elif usage_percent >= 80: + return QuotaWarningLevel.WARNING + else: + return QuotaWarningLevel.NORMAL diff --git a/packages/ports/project_repository.py b/packages/ports/project_repository.py index eae76ed99..fc60b8452 100644 --- a/packages/ports/project_repository.py +++ b/packages/ports/project_repository.py @@ -1,13 +1,30 @@ -from __future__ import annotations - -from typing import Protocol - -from packages.domain import Project +""" +Project 仓储接口 +""" +from abc import ABC, abstractmethod +from typing import Optional +from packages.domain.entities import Project -class ProjectRepository(Protocol): - def list_by_workspace(self, workspace_id: str) -> list[Project]: - """List projects for a workspace.""" - - def create(self, project: Project) -> Project: - """Persist a project and return it.""" +class ProjectRepository(ABC): + """Project 仓储接口""" + + @abstractmethod + def save(self, project: Project) -> None: + """保存项目""" + pass + + @abstractmethod + def find_by_id(self, project_id: str) -> Optional[Project]: + """根据 ID 查找项目""" + pass + + @abstractmethod + def count_by_workspace(self, workspace_id: str) -> int: + """统计工作空间的项目数量""" + pass + + @abstractmethod + def delete(self, project_id: str) -> bool: + """删除项目""" + pass diff --git a/tests/unit/test_quota.py b/tests/unit/test_quota.py new file mode 100644 index 000000000..41650eaac --- /dev/null +++ b/tests/unit/test_quota.py @@ -0,0 +1,287 @@ +""" +配额检查服务测试 +""" +import pytest +from unittest.mock import Mock +from packages.domain.quota import ( + QuotaChecker, + QuotaWarningLevel, + get_warning_level, +) +from packages.domain.entities import Workspace + + +class TestQuotaChecker: + """配额检查器测试""" + + @pytest.fixture + def mock_workspace_repo(self): + repo = Mock() + repo.find_by_id = Mock(return_value=None) + repo.save = Mock() + return repo + + @pytest.fixture + def mock_project_repo(self): + repo = Mock() + repo.count_by_workspace = Mock(return_value=0) + return repo + + @pytest.fixture + def checker(self, mock_workspace_repo, mock_project_repo): + return QuotaChecker( + workspace_repository=mock_workspace_repo, + project_repository=mock_project_repo, + ) + + def test_check_can_create_project_within_limit( + self, + checker, + mock_workspace_repo, + mock_project_repo, + ): + """测试在配额内可以创建项目""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + subscription_plan="free", + max_projects=3, + ) + mock_workspace_repo.find_by_id.return_value = workspace + mock_project_repo.count_by_workspace.return_value = 2 + + can_create, error = checker.check_can_create_project("workspace-123") + + assert can_create is True + assert error is None + + def test_check_can_create_project_at_limit( + self, + checker, + mock_workspace_repo, + mock_project_repo, + ): + """测试达到配额限制""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + subscription_plan="free", + max_projects=3, + ) + mock_workspace_repo.find_by_id.return_value = workspace + mock_project_repo.count_by_workspace.return_value = 3 + + can_create, error = checker.check_can_create_project("workspace-123") + + assert can_create is False + assert "Project limit reached" in error + + def test_check_can_create_project_unlimited( + self, + checker, + mock_workspace_repo, + mock_project_repo, + ): + """测试无限制计划""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + subscription_plan="pro", + max_projects=999999, + ) + mock_workspace_repo.find_by_id.return_value = workspace + mock_project_repo.count_by_workspace.return_value=1000 + + can_create, error = checker.check_can_create_project("workspace-123") + + assert can_create is True + assert error is None + + def test_check_storage_available_within_limit( + self, + checker, + mock_workspace_repo, + ): + """测试存储在配额内""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + subscription_plan="free", + max_storage_gb=10, + used_storage_gb=5.0, + ) + mock_workspace_repo.find_by_id.return_value = workspace + + can_store, error = checker.check_storage_available("workspace-123", 3.0) + + assert can_store is True + assert error is None + + def test_check_storage_available_exceeded( + self, + checker, + mock_workspace_repo, + ): + """测试存储超出配额""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + subscription_plan="free", + max_storage_gb=10, + used_storage_gb=8.0, + ) + mock_workspace_repo.find_by_id.return_value = workspace + + can_store, error = checker.check_storage_available("workspace-123", 3.0) + + assert can_store is False + assert "Storage limit exceeded" in error + + def test_get_quota_status( + self, + checker, + mock_workspace_repo, + mock_project_repo, + ): + """测试获取配额状态""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + subscription_plan="free", + max_projects=3, + max_storage_gb=10, + used_storage_gb=7.5, + ) + mock_workspace_repo.find_by_id.return_value = workspace + mock_project_repo.count_by_workspace.return_value = 2 + + status = checker.get_quota_status("workspace-123") + + assert status is not None + assert status["workspace_id"] == "workspace-123" + assert status["subscription_plan"] == "free" + assert status["projects"]["used"] == 2 + assert status["projects"]["limit"] == 3 + assert status["projects"]["usage_percent"] == pytest.approx(66.67, rel=0.1) + assert status["storage"]["used_gb"] == 7.5 + assert status["storage"]["limit_gb"] == 10 + assert status["storage"]["remaining_gb"] == 2.5 + assert status["storage"]["usage_percent"] == 75.0 + + def test_get_quota_status_unlimited( + self, + checker, + mock_workspace_repo, + mock_project_repo, + ): + """测试无限制配额状态""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + subscription_plan="pro", + max_projects=999999, + max_storage_gb=100, + used_storage_gb=50.0, + ) + mock_workspace_repo.find_by_id.return_value = workspace + mock_project_repo.count_by_workspace.return_value = 1000 + + status = checker.get_quota_status("workspace-123") + + assert status["projects"]["unlimited"] is True + assert status["projects"]["usage_percent"] == 0 + + def test_update_storage_usage_increase( + self, + checker, + mock_workspace_repo, + ): + """测试增加存储使用量""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + max_storage_gb=10, + used_storage_gb=5.0, + ) + mock_workspace_repo.find_by_id.return_value = workspace + + success, error = checker.update_storage_usage("workspace-123", 2.5) + + assert success is True + assert error is None + assert workspace.used_storage_gb == 7.5 + mock_workspace_repo.save.assert_called_once() + + def test_update_storage_usage_decrease( + self, + checker, + mock_workspace_repo, + ): + """测试减少存储使用量""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + max_storage_gb=10, + used_storage_gb=5.0, + ) + mock_workspace_repo.find_by_id.return_value = workspace + + success, error = checker.update_storage_usage("workspace-123", -2.0) + + assert success is True + assert workspace.used_storage_gb == 3.0 + + def test_update_storage_usage_prevent_negative( + self, + checker, + mock_workspace_repo, + ): + """测试防止负数存储""" + workspace = Workspace( + id="workspace-123", + name="Test", + owner_user_id="owner-id", + max_storage_gb=10, + used_storage_gb=2.0, + ) + mock_workspace_repo.find_by_id.return_value = workspace + + success, error = checker.update_storage_usage("workspace-123", -5.0) + + assert success is True + assert workspace.used_storage_gb == 0.0 + + +class TestWarningLevel: + """警告级别测试""" + + def test_get_warning_level_normal(self): + """测试正常级别""" + assert get_warning_level(50.0) == QuotaWarningLevel.NORMAL + assert get_warning_level(79.9) == QuotaWarningLevel.NORMAL + + def test_get_warning_level_warning(self): + """测试警告级别""" + assert get_warning_level(80.0) == QuotaWarningLevel.WARNING + assert get_warning_level(85.0) == QuotaWarningLevel.WARNING + assert get_warning_level(89.9) == QuotaWarningLevel.WARNING + + def test_get_warning_level_critical(self): + """测试严重级别""" + assert get_warning_level(90.0) == QuotaWarningLevel.CRITICAL + assert get_warning_level(95.0) == QuotaWarningLevel.CRITICAL + assert get_warning_level(99.9) == QuotaWarningLevel.CRITICAL + + def test_get_warning_level_exceeded(self): + """测试超出级别""" + assert get_warning_level(100.0) == QuotaWarningLevel.EXCEEDED + assert get_warning_level(105.0) == QuotaWarningLevel.EXCEEDED