yingqing/core/backend/apps/billing/tests.py

from decimal import Decimal

from django.test import TestCase
from rest_framework.test import APIClient

from apps.accounts.models import Team, TeamMember, User
from apps.ai.models import AITask, ModelConfig, ModelProvider
from apps.billing.models import CreditAccount, CreditLedger, CreditReservation
from apps.billing.services.ledger import charge_reserved_credit, release_credit, reserve_credit


class CreditLedgerTests(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(username="owner", password="pass")
        self.team = Team.objects.create(name="Billing Team", owner=self.user)
        TeamMember.objects.create(team=self.team, user=self.user, role=TeamMember.Role.OWNER)
        self.account = CreditAccount.objects.create(team=self.team, balance=Decimal("100.0000"))
        self.provider = ModelProvider.objects.create(name="volcengine", display_name="Volcano")
        self.model = ModelConfig.objects.create(
            provider=self.provider,
            name="doubao-seed-2-0-pro-260215",
            display_name="Doubao",
            capability=ModelConfig.Capability.TEXT,
        )
        self.task = AITask.objects.create(
            team=self.team,
            created_by=self.user,
            task_type=AITask.Type.SCRIPT_GENERATION,
            model_config=self.model,
            idempotency_key="billing-test-task",
            estimated_cost=Decimal("10.0000"),
        )

    def test_reserve_and_charge_credit(self):
        reservation = reserve_credit(team=self.team, user=self.user, task=self.task, amount=Decimal("10.0000"))
        self.account.refresh_from_db()

        self.assertEqual(reservation.status, CreditReservation.Status.ACTIVE)
        self.assertEqual(self.account.balance, Decimal("100.0000"))
        self.assertEqual(self.account.reserved_balance, Decimal("10.0000"))

        charge_reserved_credit(reservation=reservation, actual_amount=Decimal("8.0000"))
        self.account.refresh_from_db()
        reservation.refresh_from_db()

        self.assertEqual(reservation.status, CreditReservation.Status.CHARGED)
        self.assertEqual(self.account.balance, Decimal("92.0000"))
        self.assertEqual(self.account.reserved_balance, Decimal("0.0000"))
        self.assertEqual(CreditLedger.objects.filter(team=self.team).count(), 3)

    def test_release_reserved_credit(self):
        reservation = reserve_credit(team=self.team, user=self.user, task=self.task, amount=Decimal("10.0000"))
        release_credit(reservation=reservation, reason="model failed")
        self.account.refresh_from_db()
        reservation.refresh_from_db()

        self.assertEqual(reservation.status, CreditReservation.Status.RELEASED)
        self.assertEqual(self.account.balance, Decimal("100.0000"))
        self.assertEqual(self.account.reserved_balance, Decimal("0.0000"))
        self.assertEqual(CreditLedger.objects.filter(ledger_type=CreditLedger.Type.RELEASE).count(), 1)


class MemberLimitTests(TestCase):
    """成员三档额度(每日/每月/累计总额,0=不限)在 reserve_credit 时逐档管控(PMC#17/#18)。"""

    def setUp(self):
        self.user = User.objects.create_user(username="lim-owner", password="pass12345")
        self.team = Team.objects.create(name="Limit Team", owner=self.user)
        self.member = TeamMember.objects.create(team=self.team, user=self.user, role=TeamMember.Role.OWNER)
        CreditAccount.objects.create(team=self.team, balance=Decimal("1000.0000"))
        self.provider = ModelProvider.objects.create(name="volcengine", display_name="Volcano")
        self.model = ModelConfig.objects.create(
            provider=self.provider, name="m", display_name="M", capability=ModelConfig.Capability.IMAGE,
        )

    def _task(self, key):
        return AITask.objects.create(
            team=self.team, created_by=self.user, task_type=AITask.Type.PRODUCT_IMAGE,
            model_config=self.model, idempotency_key=key, estimated_cost=Decimal("10.0000"),
        )

    def test_daily_limit_blocks_over_spend(self):
        self.member.daily_credit_limit = Decimal("15.0000")
        self.member.save(update_fields=["daily_credit_limit"])
        # 第一笔 10 预留通过,第二笔 10 超每日 15 → 被拦
        reserve_credit(team=self.team, user=self.user, task=self._task("d1"), amount=Decimal("10.0000"))
        with self.assertRaises(ValueError):
            reserve_credit(team=self.team, user=self.user, task=self._task("d2"), amount=Decimal("10.0000"))

    def test_total_limit_blocks_over_spend(self):
        self.member.total_credit_limit = Decimal("12.0000")
        self.member.save(update_fields=["total_credit_limit"])
        reserve_credit(team=self.team, user=self.user, task=self._task("t1"), amount=Decimal("10.0000"))
        with self.assertRaises(ValueError):
            reserve_credit(team=self.team, user=self.user, task=self._task("t2"), amount=Decimal("10.0000"))

    def test_zero_limit_means_unlimited(self):
        # 三档都 0 = 不限:多笔大额预留都不被额度档拦(只受账户余额约束)
        for i in range(5):
            reserve_credit(team=self.team, user=self.user, task=self._task(f"z{i}"), amount=Decimal("10.0000"))


class RechargePermissionTests(TestCase):
    """充值是团队资金操作:仅 owner/admin 可发起,普通成员/访客被 403 拦截。"""

    def setUp(self):
        self.owner = User.objects.create_user(username="r-owner", password="ownerpass123")
        self.member = User.objects.create_user(username="r-member", password="memberpass123")
        self.team = Team.objects.create(name="Recharge Team", owner=self.owner)
        TeamMember.objects.create(team=self.team, user=self.owner, role=TeamMember.Role.OWNER, status=TeamMember.Status.ACTIVE)
        TeamMember.objects.create(team=self.team, user=self.member, role=TeamMember.Role.MEMBER, status=TeamMember.Status.ACTIVE)
        CreditAccount.objects.create(team=self.team, balance=Decimal("0.0000"))

    def test_owner_can_recharge(self):
        client = APIClient()
        client.force_authenticate(self.owner)
        response = client.post("/api/billing/recharge/", {"amount": "100"}, format="json")
        self.assertEqual(response.status_code, 201)
        account = CreditAccount.objects.get(team=self.team)
        self.assertEqual(account.balance, Decimal("100.0000"))

    def test_member_cannot_recharge(self):
        client = APIClient()
        client.force_authenticate(self.member)
        response = client.post("/api/billing/recharge/", {"amount": "100"}, format="json")
        self.assertEqual(response.status_code, 403)
        account = CreditAccount.objects.get(team=self.team)
        self.assertEqual(account.balance, Decimal("0.0000"))  # 余额未变,越权被拦