chunyu
008b82055d
新应用aitool：POST /api/aitool/upscale（登录扣2积分、事务行锁防并发超扣，失败退款+failed记录）；产物存MEDIA_ROOT/aitool/upscale，返回可下载URL；provider默认本地Pillow，第三方网关走环境变量不落库。
90 lines
3.9 KiB
Python
90 lines
3.9 KiB
Python
"""C-04 图片放大最小闭环用例：成功链路/余额不足/坏图回滚/流水/并发。"""
|
|
import io
|
|
|
|
from django.core.files.uploadedfile import SimpleUploadedFile
|
|
from django.test import TestCase
|
|
from PIL import Image
|
|
from rest_framework.test import APIClient
|
|
|
|
from aitool.models import AIToolRecord
|
|
from aitool.views.upscale_view import UPSCALE_COST
|
|
from user.models import FUser, PointTransaction
|
|
|
|
|
|
def _resp_bytes(response):
|
|
return b''.join(response.streaming_content)
|
|
|
|
|
|
def _img(color='red', size=(32, 32)):
|
|
img = Image.new('RGB', size, color=color)
|
|
buf = io.BytesIO()
|
|
img.save(buf, format='PNG')
|
|
buf.seek(0)
|
|
return SimpleUploadedFile('t.png', buf.getvalue(), content_type='image/png')
|
|
|
|
|
|
class UpscaleTest(TestCase):
|
|
def setUp(self):
|
|
self.client = APIClient()
|
|
self.u = FUser.objects.create_user(username='c04u', password='x12345678')
|
|
self.u.points = 100
|
|
self.u.save(update_fields=['points'])
|
|
self.client.force_authenticate(user=self.u)
|
|
|
|
def test_success_chain(self):
|
|
r = self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 2}, format='multipart')
|
|
self.assertEqual(r.status_code, 200)
|
|
self.assertEqual(r['X-AITool-Cost'], str(UPSCALE_COST))
|
|
self.assertTrue(r['X-AITool-File'].startswith('/media/aitool/upscale/'))
|
|
out = Image.open(io.BytesIO(_resp_bytes(r)))
|
|
self.assertEqual(out.size, (64, 64))
|
|
self.u.refresh_from_db()
|
|
self.assertEqual(self.u.points, 100 - UPSCALE_COST)
|
|
rec = AIToolRecord.objects.filter(user=self.u, status='success').order_by('-id').first()
|
|
self.assertIsNotNone(rec)
|
|
self.assertEqual(rec.scale, 2)
|
|
tx = PointTransaction.objects.filter(user=self.u, transaction_type='spend').order_by('-id').first()
|
|
self.assertEqual(tx.amount, UPSCALE_COST)
|
|
|
|
def test_scale_4x(self):
|
|
r = self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 4}, format='multipart')
|
|
self.assertEqual(r.status_code, 200)
|
|
out = Image.open(io.BytesIO(_resp_bytes(r)))
|
|
self.assertEqual(out.size, (128, 128))
|
|
|
|
def test_insufficient_points(self):
|
|
self.u.points = 0
|
|
self.u.save(update_fields=['points'])
|
|
r = self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 2}, format='multipart')
|
|
self.assertEqual(r.status_code, 402)
|
|
self.assertEqual(r.json().get('code'), 'INSUFFICIENT_POINTS')
|
|
self.assertFalse(AIToolRecord.objects.filter(user=self.u).exists())
|
|
|
|
def test_invalid_scale(self):
|
|
r = self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 8}, format='multipart')
|
|
self.assertEqual(r.status_code, 400)
|
|
|
|
def test_bad_image_refunds(self):
|
|
bad = SimpleUploadedFile('t.png', b'not-an-image', content_type='image/png')
|
|
before = self.u.points
|
|
r = self.client.post('/api/aitool/upscale/', {'file': bad, 'scale': 2}, format='multipart')
|
|
self.assertIn(r.status_code, (400, 500))
|
|
self.u.refresh_from_db()
|
|
self.assertEqual(self.u.points, before)
|
|
statuses = set(AIToolRecord.objects.filter(user=self.u).values_list('status', flat=True))
|
|
self.assertTrue(statuses <= {'failed', 'refunded', 'success'})
|
|
# 坏图应有回滚流水
|
|
self.assertTrue(PointTransaction.objects.filter(
|
|
user=self.u, transaction_type='earn', description__contains='回滚').exists())
|
|
|
|
def test_records_endpoint(self):
|
|
self.client.post('/api/aitool/upscale/', {'file': _img(), 'scale': 2}, format='multipart')
|
|
r = self.client.get('/api/aitool/records/')
|
|
self.assertEqual(r.status_code, 200)
|
|
self.assertTrue(len(r.json()['data']) >= 1)
|
|
|
|
def test_requires_login(self):
|
|
anon = APIClient()
|
|
r = anon.post('/api/aitool/upscale/', {'file': _img()}, format='multipart')
|
|
self.assertIn(r.status_code, (401, 403))
|