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")) # 余额未变,越权被拦