feat: add AirShelf core implementation

2026-06-05 10:21:40 +08:00
parent 2ba1058329
commit cfdcd84a30
252 changed files with 70828 additions and 0 deletions
@@ -0,0 +1,60 @@
+from decimal import Decimal
+
+from django.test import TestCase
+
+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)
+