feat: add balance API endpoint
This commit is contained in:
@@ -120,6 +120,77 @@ class ApiKeyAuthenticationTests(TestCase):
|
||||
self.assertEqual(response.data["error"]["code"], "unauthorized")
|
||||
|
||||
|
||||
class BalanceApiTests(TestCase):
|
||||
url = "/api/v1/balance"
|
||||
|
||||
def setUp(self):
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
self.user = get_user_model().objects.create_user(
|
||||
username=f"balance-user-{suffix}",
|
||||
email=f"balance-user-{suffix}@example.com",
|
||||
password="password",
|
||||
)
|
||||
self.api_key, self.raw_key = ApiKey.create_for_user(self.user, name="balance")
|
||||
self.client = APIClient()
|
||||
|
||||
def auth_header(self, raw_key: str | None = None) -> dict:
|
||||
return {"HTTP_AUTHORIZATION": f"Bearer {raw_key or self.raw_key}"}
|
||||
|
||||
def test_balance_returns_wallet_balance_matching_ledger_sum(self):
|
||||
UserWallet.objects.create(user=self.user, points_balance=100)
|
||||
PointsLedger.objects.create(
|
||||
user=self.user,
|
||||
change_type=PointsLedger.ChangeType.RECHARGE,
|
||||
points_delta=120,
|
||||
balance_after=120,
|
||||
ref_order_id=1,
|
||||
)
|
||||
call = CallRecord.objects.create(
|
||||
user=self.user,
|
||||
api_key=self.api_key,
|
||||
operation_type=CallRecord.OperationType.TITLE,
|
||||
alias="title-standard",
|
||||
model_used="gpt-5.5",
|
||||
points_cost=20,
|
||||
status=CallRecord.Status.SUCCESS,
|
||||
)
|
||||
PointsLedger.objects.create(
|
||||
user=self.user,
|
||||
change_type=PointsLedger.ChangeType.CONSUME,
|
||||
points_delta=-20,
|
||||
balance_after=100,
|
||||
ref_call=call,
|
||||
)
|
||||
|
||||
response = self.client.get(self.url, **self.auth_header())
|
||||
|
||||
ledger_sum = sum(
|
||||
PointsLedger.objects.filter(user=self.user).values_list(
|
||||
"points_delta",
|
||||
flat=True,
|
||||
)
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response.data["user"], self.user.username)
|
||||
self.assertEqual(response.data["points_balance"], 100)
|
||||
self.assertEqual(response.data["points_balance"], ledger_sum)
|
||||
|
||||
def test_balance_returns_zero_without_creating_missing_wallet(self):
|
||||
response = self.client.get(self.url, **self.auth_header())
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(response.data["points_balance"], 0)
|
||||
self.assertFalse(UserWallet.objects.filter(user=self.user).exists())
|
||||
|
||||
def test_balance_does_not_accept_web_session_without_api_key(self):
|
||||
self.client.force_login(self.user)
|
||||
|
||||
response = self.client.get(self.url)
|
||||
|
||||
self.assertEqual(response.status_code, 401)
|
||||
self.assertEqual(response.data["error"]["code"], "unauthorized")
|
||||
|
||||
|
||||
class FakeGenerationProvider:
|
||||
def __init__(self, *, capabilities=None):
|
||||
self._capabilities = set(capabilities or {"text", "image", "vision"})
|
||||
|
||||
+2
-1
@@ -1,8 +1,9 @@
|
||||
from django.urls import path
|
||||
|
||||
from .views import GenerateImageView, GenerateTitleView
|
||||
from .views import BalanceView, GenerateImageView, GenerateTitleView
|
||||
|
||||
urlpatterns = [
|
||||
path("v1/balance", BalanceView.as_view(), name="api-balance"),
|
||||
path("v1/generate/title", GenerateTitleView.as_view(), name="api-generate-title"),
|
||||
path("v1/generate/image", GenerateImageView.as_view(), name="api-generate-image"),
|
||||
]
|
||||
|
||||
@@ -15,6 +15,7 @@ from apps.api.serializers import (
|
||||
GenerateImageRequestSerializer,
|
||||
GenerateTitleRequestSerializer,
|
||||
)
|
||||
from apps.billing.services import get_balance_snapshot
|
||||
|
||||
|
||||
class ExternalApiView(APIView):
|
||||
@@ -64,3 +65,15 @@ class GenerateImageView(ExternalApiView):
|
||||
except ApiRequestError as exc:
|
||||
return Response(exc.as_response_data(), status=exc.http_status)
|
||||
return Response(data, status=status.HTTP_200_OK)
|
||||
|
||||
|
||||
class BalanceView(ExternalApiView):
|
||||
def get(self, request):
|
||||
balance = get_balance_snapshot(request.user)
|
||||
return Response(
|
||||
{
|
||||
"user": request.user.get_username(),
|
||||
"points_balance": balance.points_balance,
|
||||
},
|
||||
status=status.HTTP_200_OK,
|
||||
)
|
||||
|
||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
|
||||
from django.db import transaction
|
||||
from django.db.models import Sum
|
||||
|
||||
from apps.users.models import UserWallet
|
||||
|
||||
@@ -46,6 +47,12 @@ class RefundResult:
|
||||
refunded: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BalanceSnapshot:
|
||||
points_balance: int
|
||||
ledger_balance: int
|
||||
|
||||
|
||||
def _validate_positive_points(points: int) -> int:
|
||||
if isinstance(points, bool) or not isinstance(points, int) or points <= 0:
|
||||
raise ValueError("points must be a positive integer")
|
||||
@@ -57,6 +64,21 @@ def _locked_wallet_for_user(user) -> UserWallet:
|
||||
return wallet
|
||||
|
||||
|
||||
def get_balance_snapshot(user) -> BalanceSnapshot:
|
||||
points_balance = (
|
||||
UserWallet.objects.filter(user=user)
|
||||
.values_list("points_balance", flat=True)
|
||||
.first()
|
||||
)
|
||||
ledger_balance = PointsLedger.objects.filter(user=user).aggregate(
|
||||
total=Sum("points_delta")
|
||||
)["total"] or 0
|
||||
return BalanceSnapshot(
|
||||
points_balance=int(points_balance or 0),
|
||||
ledger_balance=int(ledger_balance),
|
||||
)
|
||||
|
||||
|
||||
def precharge_call(
|
||||
*,
|
||||
user,
|
||||
|
||||
Reference in New Issue
Block a user