xiaoxia-saas/tests/unit/test_quota.py

"""
配额检查服务测试
"""

from unittest.mock import Mock

import pytest

from packages.domain.entities import Workspace
from packages.domain.quota import QuotaChecker, QuotaWarningLevel, get_warning_level


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