xiaoxia-saas/tests/unit/test_quota.py

"""Quota 配额系统单测 — 全维度覆盖."""

from __future__ import annotations

import math

import pytest

from packages.domain.quota import (
    QUOTA_TIERS,
    QuotaChecker,
    QuotaCheckResult,
    QuotaDimension,
    QuotaRegistry,
    QuotaTier,
    QuotaWarningLevel,
    get_warning_level,
    quota_checker,
    quota_registry,
)

# ── 枚举与常量 ──────────────────────────────────────────────────────────────


class TestQuotaDimension:
    def test_storage_gb(self):
        assert QuotaDimension.STORAGE_GB == "storage_gb"

    def test_videos_per_month(self):
        assert QuotaDimension.VIDEOS_PER_MONTH == "videos_per_month"

    def test_max_concurrent(self):
        assert QuotaDimension.MAX_CONCURRENT == "max_concurrent"

    def test_max_templates(self):
        assert QuotaDimension.MAX_TEMPLATES == "max_templates"

    def test_ai_voice_enabled(self):
        assert QuotaDimension.AI_VOICE_ENABLED == "ai_voice_enabled"

    def test_all_dimensions_count(self):
        # 至少包含内置的几个核心维度
        dims = list(QuotaDimension)
        assert len(dims) >= 7

    def test_from_string(self):
        assert QuotaDimension("storage_gb") == QuotaDimension.STORAGE_GB


class TestQuotaWarningLevel:
    def test_normal(self):
        assert QuotaWarningLevel.NORMAL == "normal"

    def test_warning(self):
        assert QuotaWarningLevel.WARNING == "warning"

    def test_critical(self):
        assert QuotaWarningLevel.CRITICAL == "critical"

    def test_exceeded(self):
        assert QuotaWarningLevel.EXCEEDED == "exceeded"


# ── QuotaTier ──────────────────────────────────────────────────────────────


class TestQuotaTier:
    def test_get_limit_defined(self):
        tier = QuotaTier(name="test", limits={"storage_gb": 100})
        assert tier.get_limit("storage_gb") == 100

    def test_get_limit_undefined_returns_zero(self):
        tier = QuotaTier(name="test", limits={})
        assert tier.get_limit("unknown_dim") == 0

    def test_is_unlimited_true(self):
        tier = QuotaTier(name="test", limits={"templates": float("inf")})
        assert tier.is_unlimited("templates") is True

    def test_is_unlimited_false(self):
        tier = QuotaTier(name="test", limits={"storage_gb": 100})
        assert tier.is_unlimited("storage_gb") is False

    def test_is_unlimited_undefined_defaults_true(self):
        # 未定义的维度，get 默认为 inf → is_unlimited 返回 True
        tier = QuotaTier(name="test", limits={})
        assert tier.is_unlimited("unknown") is True


class TestQuotaTiers:
    def test_free_tier_exists(self):
        assert "free" in QUOTA_TIERS
        free = QUOTA_TIERS["free"]
        assert free.name == "free"
        assert free.get_limit("storage_gb") == 2
        assert free.get_limit("videos_per_month") == 5

    def test_basic_tier_exists(self):
        assert "basic" in QUOTA_TIERS
        basic = QUOTA_TIERS["basic"]
        assert basic.get_limit("storage_gb") == 20
        assert basic.get_limit("videos_per_month") == 30
        assert basic.get_limit("ai_voice_enabled") == 1

    def test_premium_tier_exists(self):
        assert "premium" in QUOTA_TIERS
        premium = QUOTA_TIERS["premium"]
        assert premium.get_limit("storage_gb") == 100
        assert premium.get_limit("videos_per_month") == 100

    def test_premium_templates_unlimited(self):
        premium = QUOTA_TIERS["premium"]
        assert premium.is_unlimited("max_templates") is True

    def test_free_ai_voice_disabled(self):
        free = QUOTA_TIERS["free"]
        assert free.get_limit("ai_voice_enabled") == 0

    def test_basic_ai_voice_enabled(self):
        basic = QUOTA_TIERS["basic"]
        assert basic.get_limit("ai_voice_enabled") == 1

    def test_storage_increases_with_tier(self):
        free = QUOTA_TIERS["free"].get_limit("storage_gb")
        basic = QUOTA_TIERS["basic"].get_limit("storage_gb")
        premium = QUOTA_TIERS["premium"].get_limit("storage_gb")
        assert free < basic < premium


# ── QuotaCheckResult ───────────────────────────────────────────────────────


class TestQuotaCheckResult:
    def test_usage_percent_normal(self):
        result = QuotaCheckResult(
            allowed=True,
            dimension="storage_gb",
            limit=100,
            used=50,
            remaining=50,
            warning_level=QuotaWarningLevel.NORMAL,
        )
        assert result.usage_percent == 50.0

    def test_usage_percent_zero_usage(self):
        result = QuotaCheckResult(
            allowed=True,
            dimension="storage_gb",
            limit=100,
            used=0,
            remaining=100,
            warning_level=QuotaWarningLevel.NORMAL,
        )
        assert result.usage_percent == 0.0

    def test_usage_percent_exceeded_capped(self):
        result = QuotaCheckResult(
            allowed=False,
            dimension="storage_gb",
            limit=100,
            used=150,
            remaining=0,
            warning_level=QuotaWarningLevel.EXCEEDED,
        )
        assert result.usage_percent == 100.0  # capped at 100

    def test_usage_percent_zero_limit_with_usage(self):
        result = QuotaCheckResult(
            allowed=False,
            dimension="storage_gb",
            limit=0,
            used=10,
            remaining=0,
            warning_level=QuotaWarningLevel.EXCEEDED,
        )
        assert result.usage_percent == 100.0

    def test_usage_percent_zero_limit_no_usage(self):
        result = QuotaCheckResult(
            allowed=True,
            dimension="storage_gb",
            limit=0,
            used=0,
            remaining=0,
            warning_level=QuotaWarningLevel.NORMAL,
        )
        assert result.usage_percent == 0.0

    def test_usage_percent_unlimited(self):
        result = QuotaCheckResult(
            allowed=True,
            dimension="max_templates",
            limit=float("inf"),
            used=100,
            remaining=float("inf"),
            warning_level=QuotaWarningLevel.NORMAL,
        )
        assert result.usage_percent == 0.0


# ── QuotaRegistry ──────────────────────────────────────────────────────────


class TestQuotaRegistry:
    def test_init_has_builtin_dimensions(self):
        reg = QuotaRegistry()
        dims = reg.list_dimensions()
        assert "storage_gb" in dims
        assert "videos_per_month" in dims
        assert "max_concurrent" in dims

    def test_init_has_three_tiers(self):
        reg = QuotaRegistry()
        tiers = reg.list_tiers()
        assert "free" in tiers
        assert "basic" in tiers
        assert "premium" in tiers

    def test_get_limit_free_storage(self):
        reg = QuotaRegistry()
        assert reg.get_limit("free", "storage_gb") == 2

    def test_get_limit_unknown_plan_returns_zero(self):
        reg = QuotaRegistry()
        assert reg.get_limit("unknown_plan", "storage_gb") == 2  # fallback to free

    def test_get_tier_free(self):
        reg = QuotaRegistry()
        tier = reg.get_tier("free")
        assert tier is not None
        assert tier.name == "free"

    def test_get_tier_unknown_returns_none(self):
        """未知套餐 fallback 到 free"""
        reg = QuotaRegistry()
        tier = reg.get_tier("nonexistent")
        assert tier is not None
        assert tier.name == "free"

    def test_register_new_dimension(self):
        reg = QuotaRegistry()
        reg.register_dimension("custom_dim", "自定义维度", default_limits={"free": 5, "basic": 20})
        assert "custom_dim" in reg.list_dimensions()
        assert reg.get_limit("free", "custom_dim") == 5
        assert reg.get_limit("basic", "custom_dim") == 20

    def test_register_dimension_idempotent(self):
        reg = QuotaRegistry()
        reg.register_dimension("custom_dim", "v1", default_limits={"free": 5})
        reg.register_dimension("custom_dim", "v2", default_limits={"free": 99})
        # 幂等：第二次注册不改变
        assert reg.list_dimensions()["custom_dim"] == "v1"
        assert reg.get_limit("free", "custom_dim") == 5

    def test_register_dimension_no_defaults(self):
        reg = QuotaRegistry()
        reg.register_dimension("new_dim", "新维度")
        # 默认所有套餐都是 0
        assert reg.get_limit("free", "new_dim") == 0
        assert reg.get_limit("basic", "new_dim") == 0
        assert reg.get_limit("premium", "new_dim") == 0

    def test_list_dimensions_returns_copy(self):
        reg = QuotaRegistry()
        dims = reg.list_dimensions()
        dims["fake"] = "test"
        # 修改返回值不影响内部
        assert "fake" not in reg.list_dimensions()


# ── QuotaChecker ───────────────────────────────────────────────────────────


class TestQuotaChecker:
    def test_check_within_limit(self):
        checker = QuotaChecker()
        result = checker.check("free", "storage_gb", 1.0)
        assert result.allowed is True
        assert result.limit == 2
        assert result.used == 1.0
        assert result.remaining == 1.0
        assert result.dimension == "storage_gb"

    def test_check_exceeds_limit(self):
        checker = QuotaChecker()
        result = checker.check("free", "storage_gb", 3.0)
        assert result.allowed is False
        assert result.remaining == 0

    def test_check_exactly_at_limit(self):
        # used == limit 时 allowed 为 False（必须严格小于）
        checker = QuotaChecker()
        result = checker.check("free", "storage_gb", 2.0)
        assert result.allowed is False

    def test_check_unlimited(self):
        checker = QuotaChecker()
        result = checker.check("premium", "max_templates", 1000.0)
        assert result.allowed is True
        assert math.isinf(result.remaining)
        assert result.warning_level == QuotaWarningLevel.NORMAL

    def test_check_unknown_plan(self):
        checker = QuotaChecker()
        result = checker.check("unknown", "storage_gb", 1.0)
        assert result.allowed is True  # fallback to free, limit=2, used=1.0 < 2
        assert result.limit == 2  # fallback to free tier

    def test_check_multiple(self):
        checker = QuotaChecker()
        results = checker.check_multiple(
            "free",
            {"storage_gb": 1.0, "videos_per_month": 2},
        )
        assert len(results) == 2
        assert results[0].dimension == "storage_gb"
        assert results[1].dimension == "videos_per_month"
        assert all(r.allowed for r in results)

    def test_warning_level_normal(self):
        checker = QuotaChecker()
        result = checker.check("free", "storage_gb", 1.0)  # 50%
        assert result.warning_level == QuotaWarningLevel.NORMAL

    def test_warning_level_warning(self):
        checker = QuotaChecker()
        # 80% < 95% → warning
        result = checker.check("free", "storage_gb", 1.7)  # 85%
        assert result.warning_level == QuotaWarningLevel.WARNING

    def test_warning_level_critical(self):
        checker = QuotaChecker()
        # 95% <= < 100% → critical
        result = checker.check("free", "storage_gb", 1.95)  # 97.5%
        assert result.warning_level == QuotaWarningLevel.CRITICAL

    def test_warning_level_exceeded(self):
        checker = QuotaChecker()
        result = checker.check("free", "storage_gb", 2.5)  # 125%
        assert result.warning_level == QuotaWarningLevel.EXCEEDED

    def test_warning_level_zero_limit_with_usage(self):
        checker = QuotaChecker()
        result = checker.check("free", "ai_voice_enabled", 1)  # limit=0, used=1
        assert result.warning_level == QuotaWarningLevel.EXCEEDED

    def test_warning_level_zero_limit_no_usage(self):
        checker = QuotaChecker()
        # limit=0, used=0 → 特殊处理为 normal
        # 但 allowed 是 False（0 < 0 不成立）
        result = checker.check("free", "ai_voice_enabled", 0)
        # 0 < 0 是 False → not allowed
        assert result.allowed is False

    def test_checker_uses_provided_registry(self):
        reg = QuotaRegistry()
        reg.register_dimension("custom", "自定义", default_limits={"free": 42})
        checker = QuotaChecker(reg)
        result = checker.check("free", "custom", 10)
        assert result.limit == 42
        assert result.allowed is True


# ── get_warning_level 便捷函数 ────────────────────────────────────────────


class TestGetWarningLevel:
    def test_normal_low_usage(self):
        assert get_warning_level(50, 100) == QuotaWarningLevel.NORMAL

    def test_normal_zero_usage(self):
        assert get_warning_level(0, 100) == QuotaWarningLevel.NORMAL

    def test_warning_threshold(self):
        assert get_warning_level(80, 100) == QuotaWarningLevel.WARNING

    def test_warning_between_80_and_95(self):
        assert get_warning_level(90, 100) == QuotaWarningLevel.WARNING

    def test_critical_threshold(self):
        assert get_warning_level(95, 100) == QuotaWarningLevel.CRITICAL

    def test_critical_between_95_and_100(self):
        assert get_warning_level(99, 100) == QuotaWarningLevel.CRITICAL

    def test_exceeded_at_100(self):
        assert get_warning_level(100, 100) == QuotaWarningLevel.EXCEEDED

    def test_exceeded_over_100(self):
        assert get_warning_level(150, 100) == QuotaWarningLevel.EXCEEDED

    def test_unlimited_always_normal(self):
        assert get_warning_level(9999, float("inf")) == QuotaWarningLevel.NORMAL

    def test_zero_limit_with_usage_exceeded(self):
        assert get_warning_level(1, 0) == QuotaWarningLevel.EXCEEDED

    def test_zero_limit_no_usage_normal(self):
        assert get_warning_level(0, 0) == QuotaWarningLevel.NORMAL


# ── 全局单例 ───────────────────────────────────────────────────────────────


class TestGlobalSingletons:
    def test_quota_registry_exists(self):
        assert quota_registry is not None
        assert isinstance(quota_registry, QuotaRegistry)

    def test_quota_checker_exists(self):
        assert quota_checker is not None
        assert isinstance(quota_checker, QuotaChecker)

    def test_global_checker_works(self):
        result = quota_checker.check("free", "storage_gb", 1.0)
        assert result.allowed is True
        assert result.limit == 2


class TestProTier:
    """Pro 套餐专项测试"""

    def test_pro_tier_exists(self):
        """pro 套餐存在于 QUOTA_TIERS"""
        from packages.domain.quota import QUOTA_TIERS

        assert "pro" in QUOTA_TIERS

    def test_pro_tier_same_as_premium(self):
        """pro 套餐配额与 premium 完全一致"""
        from packages.domain.quota import QUOTA_TIERS

        pro = QUOTA_TIERS["pro"]
        premium = QUOTA_TIERS["premium"]
        assert pro.limits == premium.limits

    def test_pro_tier_get_limit(self):
        """pro 套餐各维度配额正确"""
        reg = QuotaRegistry()
        assert reg.get_limit("pro", "storage_gb") == 100
        assert reg.get_limit("pro", "videos_per_month") == 100
        assert reg.get_limit("pro", "max_concurrent") == 20
        assert reg.get_limit("pro", "max_titles") == 500
        assert reg.get_limit("pro", "ai_voice_enabled") == 1