124 lines
5.4 KiB
Python
124 lines
5.4 KiB
Python
from decimal import Decimal
|
|
from unittest.mock import patch
|
|
|
|
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.assets.models import Asset
|
|
from apps.billing.models import CreditAccount, CreditLedger
|
|
from apps.ops.models import Notification
|
|
from apps.ops.views import NotificationViewSet, ensure_team_notifications
|
|
from apps.products.models import Product
|
|
from apps.projects.models import Project
|
|
|
|
|
|
class BillingNotificationTests(TestCase):
|
|
def setUp(self):
|
|
self.user = User.objects.create_user(username="owner", password="pass")
|
|
self.team = Team.objects.create(name="T", owner=self.user)
|
|
TeamMember.objects.create(team=self.team, user=self.user, role=TeamMember.Role.OWNER)
|
|
CreditAccount.objects.create(team=self.team, balance="500.0000")
|
|
self.provider = ModelProvider.objects.create(name="ops-test-prov", display_name="TSR")
|
|
self.model = ModelConfig.objects.create(
|
|
provider=self.provider, name="ops-test-image", capability=ModelConfig.Capability.IMAGE
|
|
)
|
|
self.product = Product.objects.create(team=self.team, created_by=self.user, title="补水面膜")
|
|
self.project = Project.objects.create(
|
|
team=self.team, name="补水面膜·短视频", product=self.product, created_by=self.user
|
|
)
|
|
|
|
def _charge(self, task_type, amount, balance_after):
|
|
task = AITask.objects.create(
|
|
team=self.team,
|
|
project=self.project,
|
|
task_type=task_type,
|
|
model_config=self.model,
|
|
idempotency_key=f"k-{task_type}-{amount}",
|
|
actual_cost=Decimal(amount),
|
|
)
|
|
CreditLedger.objects.create(
|
|
team=self.team,
|
|
user=self.user,
|
|
project=self.project,
|
|
task=task,
|
|
ledger_type=CreditLedger.Type.CHARGE,
|
|
amount=Decimal(amount),
|
|
balance_after=Decimal(balance_after),
|
|
reason="AI 任务扣费",
|
|
)
|
|
return task
|
|
|
|
def test_charge_becomes_billing_notification_with_real_amount(self):
|
|
self._charge("product_image", "0.32", "499.68")
|
|
self._charge("video_segment", "18.40", "481.28")
|
|
|
|
ensure_team_notifications(self.team, self.user)
|
|
|
|
billing = Notification.objects.filter(team=self.team, notification_type=Notification.Type.BILLING)
|
|
# 每笔扣费一条计费消息,且费用不为空、为真实金额
|
|
costs = sorted(n.cost_label for n in billing)
|
|
self.assertIn("¥0.32", costs)
|
|
self.assertIn("¥18.40", costs)
|
|
self.assertFalse(billing.filter(cost_label__in=["", "-"]).exists())
|
|
self.assertTrue(billing.filter(title__contains="图生成").exists())
|
|
self.assertTrue(billing.filter(title__contains="视频生成").exists())
|
|
|
|
def test_project_notification_shows_aggregated_spend(self):
|
|
self._charge("product_image", "0.32", "499.68")
|
|
self._charge("video_segment", "18.40", "481.28")
|
|
|
|
ensure_team_notifications(self.team, self.user)
|
|
|
|
proj_note = Notification.objects.get(
|
|
team=self.team, notification_type=Notification.Type.TASK, title__contains="补水面膜"
|
|
)
|
|
self.assertEqual(proj_note.cost_label, "¥18.72") # 0.32 + 18.40
|
|
|
|
def test_existing_dash_cost_is_backfilled(self):
|
|
# 早期建的项目通知费用是占位「-」;补上扣费后再跑 ensure 应回填真实金额
|
|
ensure_team_notifications(self.team, self.user)
|
|
note = Notification.objects.get(team=self.team, title__contains="补水面膜")
|
|
self.assertEqual(note.cost_label, "-")
|
|
|
|
self._charge("product_image", "0.32", "499.68")
|
|
ensure_team_notifications(self.team, self.user)
|
|
note.refresh_from_db()
|
|
self.assertEqual(note.cost_label, "¥0.32")
|
|
|
|
|
|
class NotificationMuteFilterTests(TestCase):
|
|
def setUp(self):
|
|
self.user = User.objects.create_user(username="owner", password="pass")
|
|
self.team = Team.objects.create(name="T", owner=self.user)
|
|
TeamMember.objects.create(team=self.team, user=self.user, role=TeamMember.Role.OWNER)
|
|
self.client = APIClient()
|
|
self.client.force_authenticate(self.user)
|
|
for kind, title in (
|
|
(Notification.Type.BILLING, "计费 1"),
|
|
(Notification.Type.BILLING, "计费 2"),
|
|
(Notification.Type.TASK, "任务 1"),
|
|
(Notification.Type.SYSTEM, "系统 1"),
|
|
):
|
|
Notification.objects.create(
|
|
team=self.team,
|
|
recipient=self.user,
|
|
notification_type=kind,
|
|
title=title,
|
|
dedupe_key=f"mute-test:{title}",
|
|
)
|
|
|
|
@patch.object(NotificationViewSet, "_refresh_notifications")
|
|
def test_exclude_type_hides_muted_category(self, _refresh):
|
|
res = self.client.get("/api/ops/notifications/", {"exclude_type": "billing", "page_size": 50})
|
|
self.assertEqual(res.status_code, 200)
|
|
titles = [row["title"] for row in res.data["results"]]
|
|
self.assertNotIn("计费 1", titles)
|
|
self.assertNotIn("计费 2", titles)
|
|
self.assertIn("任务 1", titles)
|
|
self.assertIn("系统 1", titles)
|
|
# chip 计数仍是全量,不被静音筛掉
|
|
self.assertEqual(res.data["type_counts"]["billing"], 2)
|
|
self.assertEqual(res.data["type_counts"]["task"], 1)
|