meida_backend/tests/unit/test_tiered_budget.py

"""四档分级降级（架构 §13.3 / T-P1-25b）的判定逻辑单测。

不连网不连库：计数用 InMemoryCostCounter，时间用注入 clock。
档位语义：阈值相对 LLM_DAILY_BUDGET_CENTS；高档位**追加**拒绝低档位名单；
critical 全档豁免；手动 override 优先于成本水位。
"""

from __future__ import annotations

from datetime import UTC, datetime
from decimal import Decimal

import pytest

from packages.llm.budget import (
    BLOCK_ALL,
    OVERRIDE_KEY,
    BudgetTier,
    DailyBudgetGuard,
    InMemoryCostCounter,
)
from packages.llm.errors import DailyBudgetExceededError, SceneDegradedError


def make_guard(
    budget: str = "100",
    *,
    enabled: bool = True,
    tiers: list[dict[str, object]] | None = None,
) -> tuple[DailyBudgetGuard, InMemoryCostCounter]:
    counter = InMemoryCostCounter()
    guard = DailyBudgetGuard(
        counter,
        daily_budget_cents=Decimal(budget),
        enabled=enabled,
        clock=lambda: datetime(2026, 8, 14, 3, 30, tzinfo=UTC),  # UTC+8 已是 8-14
        tiers=tiers,
    )
    return guard, counter


async def guard_spent(spent: str) -> DailyBudgetGuard:
    guard, _ = make_guard()
    await guard.add_cost(Decimal(spent))
    return guard


# --------------------------------------------------------------------------
# 档位边界
# --------------------------------------------------------------------------


@pytest.mark.parametrize(
    ("spent", "expected_tier"),
    [
        ("59.9", 0),
        ("60", 1),
        ("79.9", 1),
        ("80", 2),
        ("89.9", 2),
        ("90", 3),
        ("99.9", 3),
        ("100", 4),
    ],
)
async def test_tier_boundaries_follow_configured_thresholds(spent: str, expected_tier: int) -> None:
    assert await (await guard_spent(spent)).current_tier() == expected_tier


async def test_tier_names_and_cumulative_blocked_scenes() -> None:
    guard = await guard_spent("60")
    assert await guard.tier_name() == "warn"
    assert await guard.blocked_scenes() == frozenset()  # warn 档无用户可感知变化

    await guard.add_cost(Decimal("20"))  # 80%
    assert await guard.tier_name() == "degrade_secondary"
    assert await guard.blocked_scenes() == frozenset({"hd_export", "report_regen"})

    await guard.add_cost(Decimal("10"))  # 90%：追加拒绝，而不是替换名单
    assert await guard.tier_name() == "degrade_generation"
    assert await guard.blocked_scenes() == frozenset(
        {"hd_export", "report_regen", "hairstyle", "tryon"}
    )

    await guard.add_cost(Decimal("10"))  # 100%
    assert await guard.tier_name() == "block_all"
    assert await guard.blocked_scenes() == BLOCK_ALL


# --------------------------------------------------------------------------
# 各档放行 / 拒绝
# --------------------------------------------------------------------------


async def test_warn_tier_blocks_nothing() -> None:
    guard = await guard_spent("60")
    for scene in ("agent_chat", "hd_export", "tryon"):
        await guard.ensure_scene_allowed(scene=scene)


async def test_tier2_blocks_secondary_scenes_but_not_plain_chat() -> None:
    guard = await guard_spent("80")
    with pytest.raises(SceneDegradedError) as excinfo:
        await guard.ensure_scene_allowed(scene="hd_export")
    assert excinfo.value.scene == "hd_export"
    assert excinfo.value.tier_name == "degrade_secondary"
    await guard.ensure_scene_allowed(scene="agent_chat")


async def test_tier3_keeps_tier2_blocks_and_adds_generation_scenes() -> None:
    guard = await guard_spent("90")
    for scene in ("hd_export", "hairstyle", "tryon"):
        with pytest.raises(SceneDegradedError):
            await guard.ensure_scene_allowed(scene=scene)
    await guard.ensure_scene_allowed(scene="agent_chat")


async def test_tier4_blocks_all_except_the_allow_list() -> None:
    guard = await guard_spent("100")
    with pytest.raises(DailyBudgetExceededError) as excinfo:
        await guard.ensure_scene_allowed(scene="tryon")
    assert excinfo.value.scene == "tryon"
    # 架构原文「仅保留在途任务与纯文本对话」：agent_chat 默认在 allow 名单
    await guard.ensure_scene_allowed(scene="agent_chat")


async def test_tier4_allow_list_is_configurable() -> None:
    tiers: list[dict[str, object]] = [
        # 裸字符串 "*" 是合法配置写法，用于演练「连纯文本对话也拦」
        {"threshold": 1.0, "name": "block_all", "blocked": BLOCK_ALL, "allow": []},
    ]
    guard, _ = make_guard(tiers=tiers)
    await guard.add_cost(Decimal("100"))
    assert await guard.current_tier() == 1  # 只有一档时 override/档位序号随之缩短
    with pytest.raises(DailyBudgetExceededError):
        await guard.ensure_scene_allowed(scene="agent_chat")


async def test_critical_scenes_are_exempt_from_every_tier() -> None:
    guard = await guard_spent("500")
    # 意图识别把守越界拒答：降级它等于让超范围提问漏过去，因此全档豁免。
    await guard.ensure_scene_allowed(scene="intent_classify", critical=True)


async def test_disabled_guard_reports_no_tier() -> None:
    guard, _ = make_guard(enabled=False)
    await guard.add_cost(Decimal("999"))
    assert await guard.current_tier() == 0
    await guard.ensure_scene_allowed(scene="tryon")


async def test_legacy_entry_point_still_enforces() -> None:
    """ensure_within_budget 兼容入口：100% 档非 allow 场景仍抛既有异常。"""
    guard = await guard_spent("100")
    with pytest.raises(DailyBudgetExceededError):
        await guard.ensure_within_budget(scene="tryon")


# --------------------------------------------------------------------------
# 手动 override
# --------------------------------------------------------------------------


async def test_override_forces_a_tier_above_the_water_level() -> None:
    guard = await guard_spent("10")  # 水位仅 10%，本未触发任何档
    await guard.set_override(4)
    assert await guard.current_tier() == 4
    with pytest.raises(DailyBudgetExceededError):
        await guard.ensure_scene_allowed(scene="tryon")
    await guard.clear_override()
    assert await guard.current_tier() == 0  # 清除后回落水位判定
    await guard.ensure_scene_allowed(scene="tryon")


async def test_override_can_loosen_a_tier_the_water_level_reached() -> None:
    """运营解除加严：水位已到 tier4，手动放宽回 tier2。"""
    guard = await guard_spent("100")
    await guard.set_override(2)
    assert await guard.current_tier() == 2
    with pytest.raises(SceneDegradedError):
        await guard.ensure_scene_allowed(scene="hd_export")
    await guard.ensure_scene_allowed(scene="tryon")


@pytest.mark.parametrize("bad_value", ["0", "5", "abc", ""])
async def test_invalid_override_values_are_ignored(bad_value: str) -> None:
    guard, counter = make_guard()
    await guard.add_cost(Decimal("100"))  # 水位 tier4
    await counter.set(OVERRIDE_KEY, bad_value, 60)
    assert await guard.current_tier() == 4  # 回落水位判定，override 不生效


async def test_set_override_rejects_out_of_range_tiers() -> None:
    guard, _ = make_guard()
    with pytest.raises(ValueError, match=r"1\.\.4"):
        await guard.set_override(5)
    with pytest.raises(ValueError, match=r"1\.\.4"):
        await guard.set_override(0)


async def test_override_ttl_expires_at_the_utc8_day_boundary() -> None:
    """override TTL 到当日 UTC+8 24 点：跨预算日必然失效，不需要人工清理。"""
    set_expiries: dict[str, int] = {}
    counter = InMemoryCostCounter()
    real_set = counter.set

    async def recording_set(key: str, value: str, ttl_seconds: int) -> None:
        set_expiries[key] = ttl_seconds
        await real_set(key, value, ttl_seconds)

    counter.set = recording_set  # type: ignore[method-assign]
    # 8-14 11:30 UTC+8 → 24:00 还有 12.5 小时
    guard = DailyBudgetGuard(
        counter,
        daily_budget_cents=Decimal("100"),
        clock=lambda: datetime(2026, 8, 14, 3, 30, tzinfo=UTC),
    )
    await guard.set_override(4)
    assert set_expiries[OVERRIDE_KEY] == 12 * 3600 + 1800


# --------------------------------------------------------------------------
# 配置解析
# --------------------------------------------------------------------------


def test_budget_tier_normalizes_bare_star_and_defaults() -> None:
    tier = BudgetTier.model_validate({"threshold": 1.0, "name": "block_all", "blocked": "*"})
    assert tier.blocked == [BLOCK_ALL]
    assert tier.blocks_all() is True
    assert tier.allow == []  # allow 未配置时默认空：block_all 档将拒绝全部非 critical
    assert BudgetTier.model_validate({"threshold": 0.6, "name": "warn", "blocked": []}).allow == []


async def test_custom_thresholds_drive_the_tiers() -> None:
    tiers: list[dict[str, object]] = [
        {"threshold": 0.5, "name": "warn", "blocked": []},
        {"threshold": 0.75, "name": "degrade_secondary", "blocked": ["hd_export"]},
        {"threshold": 1.0, "name": "block_all", "blocked": BLOCK_ALL, "allow": ["agent_chat"]},
    ]
    guard, _ = make_guard(tiers=tiers)
    await guard.add_cost(Decimal("75"))
    assert await guard.current_tier() == 2
    with pytest.raises(SceneDegradedError):
        await guard.ensure_scene_allowed(scene="hd_export")