forked from erp-dev/erp
feat: big version, added tasks for backup_database and stock change, added health check api, approve sse (support channel via merchant)
This commit is contained in:
@@ -9,7 +9,7 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
WORKDIR /app
|
||||
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends build-essential curl \
|
||||
&& apt-get install -y --no-install-recommends build-essential curl postgresql-client \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY pyproject.toml /app/
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# 保持空模块，避免在 Django 应用加载期间触发视图导入
|
||||
|
||||
|
||||
126
api_v1/tasks.py
Normal file
126
api_v1/tasks.py
Normal file
@@ -0,0 +1,126 @@
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
from celery import shared_task
|
||||
from django.conf import settings
|
||||
from django.utils import timezone
|
||||
|
||||
from basic_info import models as basic_models
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@shared_task(bind=True)
|
||||
def ping_task(self, message: str = 'ping'):
|
||||
"""
|
||||
最简单的心跳任务，用于验证 Celery worker 是否能够
|
||||
正确消费队列并返回结果。
|
||||
"""
|
||||
payload = {
|
||||
'task_id': self.request.id,
|
||||
'message': message,
|
||||
'timestamp': timezone.now().isoformat(),
|
||||
}
|
||||
logger.info('Celery ping_task 执行成功: %s', payload)
|
||||
return payload
|
||||
|
||||
|
||||
@shared_task(bind=True)
|
||||
def merchant_product_count(self, merchant_id: int):
|
||||
"""
|
||||
计算指定商户下的产品数量，用于演示如何在任务中访问数据库。
|
||||
"""
|
||||
count = basic_models.Product.objects.filter(merchant_id=merchant_id).count()
|
||||
payload = {
|
||||
'task_id': self.request.id,
|
||||
'merchant_id': merchant_id,
|
||||
'product_count': count,
|
||||
'calculated_at': timezone.now().isoformat(),
|
||||
}
|
||||
logger.info(
|
||||
'Celery merchant_product_count 统计完成: merchant=%s count=%s task=%s',
|
||||
merchant_id,
|
||||
count,
|
||||
self.request.id,
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
def _ensure_backup_dir(output_dir: str | None) -> Path:
|
||||
base_dir = Path(settings.BASE_DIR)
|
||||
backup_dir = Path(output_dir) if output_dir else (base_dir / 'data-bak')
|
||||
backup_dir.mkdir(parents=True, exist_ok=True)
|
||||
return backup_dir
|
||||
|
||||
|
||||
def _build_backup_path(backup_dir: Path, filename_prefix: str) -> Path:
|
||||
timestamp = timezone.now().strftime('%Y%m%d-%H%M%S')
|
||||
return backup_dir / f'{filename_prefix}-{timestamp}.sql'
|
||||
|
||||
|
||||
def _run_pg_dump(backup_path: Path):
|
||||
db_settings = settings.DATABASES['default']
|
||||
pg_dump = shutil.which('pg_dump')
|
||||
if not pg_dump:
|
||||
raise RuntimeError('pg_dump 不存在，请确认 PostgreSQL 客户端工具已安装')
|
||||
|
||||
host = db_settings.get('HOST') or 'localhost'
|
||||
port = db_settings.get('PORT') or '5432'
|
||||
user = db_settings.get('USER') or ''
|
||||
name = db_settings['NAME']
|
||||
password = db_settings.get('PASSWORD') or ''
|
||||
|
||||
cmd = [
|
||||
pg_dump,
|
||||
'-h',
|
||||
host,
|
||||
'-p',
|
||||
str(port),
|
||||
'-U',
|
||||
user,
|
||||
'-F',
|
||||
'p',
|
||||
'-d',
|
||||
name,
|
||||
]
|
||||
|
||||
env = os.environ.copy()
|
||||
if password:
|
||||
env['PGPASSWORD'] = password
|
||||
|
||||
with backup_path.open('wb') as stream:
|
||||
subprocess.run(cmd, check=True, stdout=stream, env=env)
|
||||
|
||||
|
||||
def _dump_database_to_sql(backup_path: Path):
|
||||
engine = settings.DATABASES['default']['ENGINE']
|
||||
if 'postgresql' not in engine:
|
||||
raise NotImplementedError('当前项目只支持 PostgreSQL 数据库备份，请检查 DATABASES 配置')
|
||||
_run_pg_dump(backup_path)
|
||||
|
||||
|
||||
@shared_task(bind=True)
|
||||
def backup_database(self, output_dir: str | None = None, filename_prefix: str = 'db-backup'):
|
||||
"""
|
||||
备份当前数据库为 .sql 文件（仅数据，不包含表结构 DDL），存放在项目根目录 data-bak 下。
|
||||
|
||||
参数:
|
||||
output_dir: 可选，指定备份目录（默认 BASE_DIR/data-bak）
|
||||
filename_prefix: 备份文件名前缀
|
||||
"""
|
||||
backup_dir = _ensure_backup_dir(output_dir)
|
||||
backup_path = _build_backup_path(backup_dir, filename_prefix)
|
||||
_dump_database_to_sql(backup_path)
|
||||
|
||||
payload = {
|
||||
'task_id': self.request.id,
|
||||
'backup_path': str(backup_path),
|
||||
'created_at': timezone.now().isoformat(),
|
||||
}
|
||||
logger.info('数据库备份完成: %s', payload)
|
||||
return payload
|
||||
|
||||
144
api_v1/tests.py
144
api_v1/tests.py
@@ -1,8 +1,25 @@
|
||||
from django.test import TestCase
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.test import TestCase, override_settings
|
||||
from django.contrib.auth.models import User, Permission
|
||||
from rest_framework.test import APIClient
|
||||
from rest_framework import status
|
||||
from basic_info.models import Merchant, UserProfile
|
||||
from basic_info.models import (
|
||||
Employee,
|
||||
Merchant,
|
||||
MerchantTypeEnum,
|
||||
Product,
|
||||
ProductCategory,
|
||||
ProductUnitEnum,
|
||||
Supplier,
|
||||
UserProfile,
|
||||
WareHouse,
|
||||
WareHouseModeEnum,
|
||||
)
|
||||
from api_v1 import tasks
|
||||
|
||||
|
||||
class UserCreationAPITestCase(TestCase):
|
||||
@@ -110,3 +127,126 @@ class UserCreationAPITestCase(TestCase):
|
||||
response = self.client.post('/api/v1/users/create/', data, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
||||
self.assertIn('password', response.data)
|
||||
|
||||
|
||||
@override_settings(
|
||||
CELERY_TASK_ALWAYS_EAGER=True,
|
||||
CELERY_TASK_EAGER_PROPAGATES=True,
|
||||
)
|
||||
class PurchaseOrderAPITestCase(TestCase):
|
||||
"""采购单 API 测试"""
|
||||
|
||||
def setUp(self):
|
||||
self.merchant = Merchant.objects.create(name='PO商户', type=MerchantTypeEnum.FACTORY)
|
||||
self.supplier = Supplier.objects.create(merchant=self.merchant, name='供应商A')
|
||||
self.warehouse = WareHouse.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='主仓',
|
||||
mode=WareHouseModeEnum.RESTRICT_IN,
|
||||
)
|
||||
category = ProductCategory.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='品类',
|
||||
product_prefix='FAB',
|
||||
)
|
||||
self.product = Product.objects.create(
|
||||
merchant=self.merchant,
|
||||
category=category,
|
||||
name='产品1',
|
||||
human_id='FAB-001',
|
||||
unit=ProductUnitEnum.METER,
|
||||
)
|
||||
self.user = User.objects.create_user(username='po_user', password='pass123')
|
||||
self.employee = Employee.objects.create(
|
||||
merchant=self.merchant,
|
||||
sys_user=self.user,
|
||||
name='仓管',
|
||||
)
|
||||
self.client = APIClient()
|
||||
self.client.force_authenticate(user=self.user)
|
||||
self.payload = {
|
||||
'supplier': self.supplier.id,
|
||||
'warehouse': self.warehouse.id,
|
||||
'order_date': '2025-11-26',
|
||||
'total_amount': '1500.00',
|
||||
'items': [
|
||||
{
|
||||
'product_id': self.product.id,
|
||||
'quantities': ['10.0', '5.0'],
|
||||
}
|
||||
],
|
||||
'remarks': '接口测试',
|
||||
}
|
||||
|
||||
def test_create_purchase_order_success(self):
|
||||
with patch('business.services.create_purchase_order_stock_entries.delay') as mock_delay:
|
||||
response = self.client.post('/api/v1/purchase-orders/', self.payload, format='json')
|
||||
|
||||
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
|
||||
self.assertIn('id', response.data)
|
||||
mock_delay.assert_called_once()
|
||||
|
||||
def test_create_purchase_order_invalid_supplier(self):
|
||||
payload = {**self.payload, 'supplier': 999}
|
||||
response = self.client.post('/api/v1/purchase-orders/', payload, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
||||
self.assertIn('不存在', response.data['error'])
|
||||
|
||||
def test_create_purchase_order_unauthenticated(self):
|
||||
self.client.force_authenticate(user=None)
|
||||
response = self.client.post('/api/v1/purchase-orders/', self.payload, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
|
||||
|
||||
|
||||
@override_settings(
|
||||
CELERY_TASK_ALWAYS_EAGER=True,
|
||||
CELERY_TASK_EAGER_PROPAGATES=True,
|
||||
)
|
||||
class CeleryTasksTestCase(TestCase):
|
||||
"""验证 Celery 任务的执行结果"""
|
||||
|
||||
def setUp(self):
|
||||
self.merchant = Merchant.objects.create(name='Celery商户', type=MerchantTypeEnum.FACTORY)
|
||||
category = ProductCategory.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='Celery品类',
|
||||
product_prefix='CLY',
|
||||
)
|
||||
for idx in range(3):
|
||||
Product.objects.create(
|
||||
merchant=self.merchant,
|
||||
category=category,
|
||||
name=f'Celery产品{idx}',
|
||||
human_id=f'CLY-{idx:03d}',
|
||||
unit=ProductUnitEnum.METER,
|
||||
)
|
||||
|
||||
def test_ping_task_returns_payload(self):
|
||||
result = tasks.ping_task.delay('celery hello')
|
||||
payload = result.get(timeout=5)
|
||||
self.assertEqual(payload['message'], 'celery hello')
|
||||
self.assertIn('timestamp', payload)
|
||||
self.assertIn('task_id', payload)
|
||||
|
||||
def test_merchant_product_count(self):
|
||||
result = tasks.merchant_product_count.delay(self.merchant.id)
|
||||
payload = result.get(timeout=5)
|
||||
self.assertEqual(payload['merchant_id'], self.merchant.id)
|
||||
self.assertEqual(payload['product_count'], 3)
|
||||
|
||||
def test_backup_database_creates_file(self):
|
||||
tmpdir = Path(tempfile.mkdtemp())
|
||||
self.addCleanup(lambda: shutil.rmtree(tmpdir, ignore_errors=True))
|
||||
with patch('api_v1.tasks._dump_database_to_sql') as mock_dump:
|
||||
def fake_dump(path: Path):
|
||||
path.write_text('-- dummy sql\n', encoding='utf-8')
|
||||
|
||||
mock_dump.side_effect = fake_dump
|
||||
result = tasks.backup_database.delay(output_dir=str(tmpdir), filename_prefix='test-backup')
|
||||
payload = result.get(timeout=10)
|
||||
backup_path = Path(payload['backup_path'])
|
||||
self.assertTrue(backup_path.exists())
|
||||
self.assertTrue(backup_path.is_file())
|
||||
self.assertEqual(backup_path.parent.resolve(), tmpdir.resolve())
|
||||
self.assertEqual(backup_path.suffix, '.sql')
|
||||
self.assertTrue(backup_path.read_text(encoding='utf-8').strip())
|
||||
|
||||
@@ -1,6 +1,15 @@
|
||||
from django.urls import path, include
|
||||
from rest_framework.routers import DefaultRouter
|
||||
from .views import stock_change_views, user_info, inventory, product_image, stateflow, users
|
||||
from .views import (
|
||||
stock_change_views,
|
||||
purchase_order,
|
||||
healthy,
|
||||
user_info,
|
||||
inventory,
|
||||
product_image,
|
||||
stateflow,
|
||||
users,
|
||||
)
|
||||
from .views.stock_change_views.snapshot import StockSnapshotListView
|
||||
from .views.printing.views import PrintingOrderViewSet, PrintingJobViewSet, PlateOrderViewSet
|
||||
from .views.upload import UploadFileViewSet
|
||||
@@ -46,6 +55,8 @@ urlpatterns = [
|
||||
|
||||
# 库存查询 API
|
||||
path('inventory/', inventory.InventoryAPIView.as_view(), name='inventory'),
|
||||
path('purchase-orders/', purchase_order.PurchaseOrderView.as_view(), name='purchase_orders'),
|
||||
path('health/', healthy.HealthCheckView.as_view(), name='health_check'),
|
||||
|
||||
# 产品图片上传 API
|
||||
path('products/<int:product_id>/image/', product_image.ProductImageUploadView.as_view(), name='product_image_upload'),
|
||||
|
||||
42
api_v1/views/healthy.py
Normal file
42
api_v1/views/healthy.py
Normal file
@@ -0,0 +1,42 @@
|
||||
import platform
|
||||
from typing import Dict
|
||||
|
||||
from django.db import connections, DatabaseError
|
||||
from django.utils import timezone
|
||||
from django.utils.version import get_version
|
||||
from rest_framework.response import Response
|
||||
from rest_framework.views import APIView
|
||||
|
||||
|
||||
class HealthCheckView(APIView):
|
||||
"""
|
||||
简单健康检查接口，返回当前服务的基础指标信息，包括：
|
||||
- 应用版本
|
||||
- Python/Django 版本
|
||||
- 服务器时间
|
||||
- 数据库连接状态
|
||||
"""
|
||||
|
||||
authentication_classes: list = []
|
||||
permission_classes: list = []
|
||||
|
||||
def get(self, request):
|
||||
db_status: Dict[str, str] = {}
|
||||
for alias in connections:
|
||||
try:
|
||||
connections[alias].cursor()
|
||||
db_status[alias] = 'ok'
|
||||
except DatabaseError as exc:
|
||||
db_status[alias] = f'error: {exc.__class__.__name__}'
|
||||
|
||||
payload = {
|
||||
'service': 'flower-api',
|
||||
'status': 'ok' if all(status == 'ok' for status in db_status.values()) else 'degraded',
|
||||
'server_time': timezone.now().isoformat(),
|
||||
'python_version': platform.python_version(),
|
||||
'django_version': get_version(),
|
||||
'platform': platform.platform(),
|
||||
'databases': db_status,
|
||||
}
|
||||
return Response(payload)
|
||||
|
||||
72
api_v1/views/purchase_order.py
Normal file
72
api_v1/views/purchase_order.py
Normal file
@@ -0,0 +1,72 @@
|
||||
from decimal import Decimal, InvalidOperation
|
||||
|
||||
from rest_framework import status, views
|
||||
from rest_framework.permissions import IsAuthenticated
|
||||
from rest_framework.response import Response
|
||||
|
||||
from basic_info import models as basic_models
|
||||
from business import services as business_services
|
||||
from .stock_change_views.mixins import StockChangeViewMixin
|
||||
|
||||
|
||||
class PurchaseOrderView(StockChangeViewMixin, views.APIView):
|
||||
"""创建采购订单并触发入库任务"""
|
||||
|
||||
permission_classes = [IsAuthenticated]
|
||||
|
||||
def post(self, request):
|
||||
if not self.check_employee_permission(request):
|
||||
return self.permission_error_response('无权限访问')
|
||||
|
||||
merchant = request.user.employee.merchant
|
||||
data = request.data or {}
|
||||
|
||||
supplier_id = data.get('supplier')
|
||||
warehouse_id = data.get('warehouse')
|
||||
order_date = data.get('order_date')
|
||||
total_amount = data.get('total_amount')
|
||||
items = data.get('items', [])
|
||||
remarks = data.get('remarks', '')
|
||||
|
||||
if not supplier_id:
|
||||
return Response({'error': '缺少供应商 ID'}, status=status.HTTP_400_BAD_REQUEST)
|
||||
if not warehouse_id:
|
||||
return Response({'error': '缺少仓库 ID'}, status=status.HTTP_400_BAD_REQUEST)
|
||||
|
||||
try:
|
||||
supplier = basic_models.Supplier.objects.get(id=supplier_id, merchant=merchant)
|
||||
except basic_models.Supplier.DoesNotExist:
|
||||
return Response({'error': f'供应商 {supplier_id} 不存在'}, status=status.HTTP_400_BAD_REQUEST)
|
||||
|
||||
try:
|
||||
basic_models.WareHouse.objects.get(id=warehouse_id, merchant=merchant)
|
||||
except basic_models.WareHouse.DoesNotExist:
|
||||
return Response({'error': f'仓库 {warehouse_id} 不存在'}, status=status.HTTP_400_BAD_REQUEST)
|
||||
|
||||
try:
|
||||
total_amount_decimal = Decimal(str(total_amount))
|
||||
except (InvalidOperation, TypeError):
|
||||
return Response({'error': 'total_amount 必须为合法数值'}, status=status.HTTP_400_BAD_REQUEST)
|
||||
|
||||
try:
|
||||
purchase_order = business_services.create_purchase_order(
|
||||
merchant=merchant,
|
||||
supplier=supplier,
|
||||
order_date=order_date,
|
||||
total_amount=total_amount_decimal,
|
||||
warehouse_id=warehouse_id,
|
||||
items=items,
|
||||
remarks=remarks,
|
||||
created_by=request.user,
|
||||
)
|
||||
except ValueError as exc:
|
||||
return Response({'error': str(exc)}, status=status.HTTP_400_BAD_REQUEST)
|
||||
|
||||
return Response(
|
||||
{
|
||||
'id': purchase_order.id,
|
||||
'message': '采购单创建成功，入库任务已排队',
|
||||
},
|
||||
status=status.HTTP_201_CREATED,
|
||||
)
|
||||
|
||||
@@ -67,6 +67,7 @@ class StockChangeViewMixin:
|
||||
"""构建明细数据"""
|
||||
return {
|
||||
'id': detail.id,
|
||||
'stock_change_record': detail.stock_change_record_id,
|
||||
'product': detail.product_id,
|
||||
'product_name': detail.product.name,
|
||||
'quantity': float(detail.quantity),
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from decimal import Decimal
|
||||
import logging
|
||||
|
||||
from django.contrib.auth import get_user_model
|
||||
from django.test import TestCase
|
||||
@@ -12,6 +13,279 @@ from stock import models as stock_models, services as stock_services
|
||||
User = get_user_model()
|
||||
|
||||
|
||||
class CreateStockChangeAPITestCase(TestCase):
|
||||
"""测试标准库存变动创建 API"""
|
||||
|
||||
def setUp(self):
|
||||
self.client = APIClient()
|
||||
|
||||
self.merchant = basic_models.Merchant.objects.create(
|
||||
name='测试商户',
|
||||
type=basic_models.MerchantTypeEnum.FACTORY,
|
||||
)
|
||||
self.category = basic_models.ProductCategory.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='布料',
|
||||
product_prefix='FAB',
|
||||
)
|
||||
self.product = basic_models.Product.objects.create(
|
||||
merchant=self.merchant,
|
||||
category=self.category,
|
||||
name='测试布料',
|
||||
human_id='FAB-001',
|
||||
unit=basic_models.ProductUnitEnum.METER,
|
||||
)
|
||||
self.warehouse = basic_models.WareHouse.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='严谨仓库',
|
||||
mode=basic_models.WareHouseModeEnum.RESTRICT_IN,
|
||||
)
|
||||
self.other_merchant = basic_models.Merchant.objects.create(
|
||||
name='其他商户',
|
||||
type=basic_models.MerchantTypeEnum.FACTORY,
|
||||
)
|
||||
other_category = basic_models.ProductCategory.objects.create(
|
||||
merchant=self.other_merchant,
|
||||
name='其他布料',
|
||||
product_prefix='FABO',
|
||||
)
|
||||
self.other_product = basic_models.Product.objects.create(
|
||||
merchant=self.other_merchant,
|
||||
category=other_category,
|
||||
name='外部布料',
|
||||
human_id='FAB-999',
|
||||
unit=basic_models.ProductUnitEnum.METER,
|
||||
)
|
||||
|
||||
self.user = User.objects.create_user(username='strict_user', password='pass123')
|
||||
self.employee = basic_models.Employee.objects.create(
|
||||
merchant=self.merchant,
|
||||
sys_user=self.user,
|
||||
name='仓管员',
|
||||
mobile='13800138002',
|
||||
status=basic_models.EmployeeStatusEnum.ACTIVE,
|
||||
)
|
||||
self.client.force_authenticate(user=self.user)
|
||||
|
||||
def test_create_stock_change_success(self):
|
||||
payload = {
|
||||
'type': stock_models.StockChangeTypeEnum.ADD,
|
||||
'warehouse': self.warehouse.id,
|
||||
'source_type': stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
'source_id': 10,
|
||||
'products': [
|
||||
{'product': self.product.id, 'quantity': ['10.50', '5.25']},
|
||||
]
|
||||
}
|
||||
|
||||
response = self.client.post('/api/v1/stock-change/', payload, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
|
||||
self.assertEqual(response.data['created_details_count'], 2)
|
||||
self.assertEqual(len(response.data['details']), 2)
|
||||
self.assertEqual(response.data['stock_change_record']['warehouse'], self.warehouse.id)
|
||||
|
||||
record_id = response.data['stock_change_record']['id']
|
||||
record = stock_models.StockChangeRecord.objects.get(id=record_id)
|
||||
self.assertEqual(record.details.count(), 2)
|
||||
|
||||
def test_create_stock_change_products_required(self):
|
||||
payload = {
|
||||
'type': stock_models.StockChangeTypeEnum.ADD,
|
||||
'warehouse': self.warehouse.id,
|
||||
'source_type': stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
'source_id': 11,
|
||||
'products': [],
|
||||
}
|
||||
|
||||
response = self.client.post('/api/v1/stock-change/', payload, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
||||
self.assertEqual(response.data['error'], '产品列表不能为空')
|
||||
|
||||
def test_create_stock_change_product_not_visible(self):
|
||||
payload = {
|
||||
'type': stock_models.StockChangeTypeEnum.ADD,
|
||||
'warehouse': self.warehouse.id,
|
||||
'source_type': stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
'products': [
|
||||
{'product': self.other_product.id, 'quantity': ['5.00']},
|
||||
]
|
||||
}
|
||||
|
||||
response = self.client.post('/api/v1/stock-change/', payload, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
|
||||
self.assertEqual(
|
||||
response.data['error'],
|
||||
f'产品ID {self.other_product.id} 对当前用户不可见',
|
||||
)
|
||||
|
||||
def test_create_stock_change_no_employee_permission_denied(self):
|
||||
class DummyUser:
|
||||
def __init__(self, username):
|
||||
self.username = username
|
||||
self.is_authenticated = True
|
||||
|
||||
dummy_user = DummyUser('no_emp_user')
|
||||
self.client.force_authenticate(user=dummy_user)
|
||||
payload = {
|
||||
'type': stock_models.StockChangeTypeEnum.ADD,
|
||||
'warehouse': self.warehouse.id,
|
||||
'source_type': stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
'products': [
|
||||
{'product': self.product.id, 'quantity': ['5.00']},
|
||||
]
|
||||
}
|
||||
|
||||
logging.disable(logging.NOTSET)
|
||||
self.addCleanup(logging.disable, logging.CRITICAL)
|
||||
with self.assertLogs('api_v1.views.stock_change_views.mixins', level='ERROR') as cm:
|
||||
response = self.client.post('/api/v1/stock-change/', payload, format='json')
|
||||
|
||||
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
|
||||
self.assertEqual(response.data['error'], '无权限访问')
|
||||
self.assertTrue(any('无员工信息' in msg for msg in cm.output))
|
||||
|
||||
|
||||
class CreateStockChangeRelaxedAPITestCase(TestCase):
|
||||
"""测试宽松模式库存变动创建 API"""
|
||||
|
||||
def setUp(self):
|
||||
self.client = APIClient()
|
||||
|
||||
self.merchant = basic_models.Merchant.objects.create(
|
||||
name='宽松商户',
|
||||
type=basic_models.MerchantTypeEnum.FACTORY,
|
||||
)
|
||||
self.category = basic_models.ProductCategory.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='面料',
|
||||
product_prefix='FAB',
|
||||
)
|
||||
self.product = basic_models.Product.objects.create(
|
||||
merchant=self.merchant,
|
||||
category=self.category,
|
||||
name='宽松布料',
|
||||
human_id='FAB-100',
|
||||
unit=basic_models.ProductUnitEnum.METER,
|
||||
)
|
||||
self.relaxed_warehouse = basic_models.WareHouse.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='宽进宽出仓',
|
||||
mode=basic_models.WareHouseModeEnum.UNRESTRICTED,
|
||||
)
|
||||
self.other_merchant = basic_models.Merchant.objects.create(
|
||||
name='RelaxedOther',
|
||||
type=basic_models.MerchantTypeEnum.FACTORY,
|
||||
)
|
||||
other_category = basic_models.ProductCategory.objects.create(
|
||||
merchant=self.other_merchant,
|
||||
name='外部面料',
|
||||
product_prefix='REL',
|
||||
)
|
||||
self.other_product = basic_models.Product.objects.create(
|
||||
merchant=self.other_merchant,
|
||||
category=other_category,
|
||||
name='外部布料',
|
||||
human_id='REL-002',
|
||||
unit=basic_models.ProductUnitEnum.METER,
|
||||
)
|
||||
|
||||
self.user = User.objects.create_user(username='relaxed_user', password='pass123')
|
||||
self.employee = basic_models.Employee.objects.create(
|
||||
merchant=self.merchant,
|
||||
sys_user=self.user,
|
||||
name='宽松仓管员',
|
||||
mobile='13800138003',
|
||||
status=basic_models.EmployeeStatusEnum.ACTIVE,
|
||||
)
|
||||
self.client.force_authenticate(user=self.user)
|
||||
|
||||
def test_create_stock_change_relaxed_success(self):
|
||||
payload = {
|
||||
'type': stock_models.StockChangeTypeEnum.ADD,
|
||||
'warehouse': self.relaxed_warehouse.id,
|
||||
'source_type': stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
'source_id': 20,
|
||||
'products': [
|
||||
{
|
||||
'product': self.product.id,
|
||||
'quantity': {'value': '12', 'unit_count': '5'}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
response = self.client.post('/api/v1/stock-change/relaxed/', payload, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
|
||||
self.assertEqual(response.data['created_details_count'], 3)
|
||||
record_id = response.data['stock_change_record']['id']
|
||||
details = stock_models.StockChangeDetail.objects.filter(stock_change_record_id=record_id)
|
||||
self.assertEqual(details.count(), 3)
|
||||
|
||||
def test_create_stock_change_relaxed_missing_required(self):
|
||||
payload = {
|
||||
'warehouse': self.relaxed_warehouse.id,
|
||||
'source_type': stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
'products': [
|
||||
{
|
||||
'product': self.product.id,
|
||||
'quantity': {'value': '10', 'unit_count': '4'}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
response = self.client.post('/api/v1/stock-change/relaxed/', payload, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
||||
self.assertEqual(response.data['error'], '缺少必要参数')
|
||||
|
||||
def test_create_stock_change_relaxed_product_not_visible(self):
|
||||
payload = {
|
||||
'type': stock_models.StockChangeTypeEnum.ADD,
|
||||
'warehouse': self.relaxed_warehouse.id,
|
||||
'source_type': stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
'products': [
|
||||
{
|
||||
'product': self.other_product.id,
|
||||
'quantity': {'value': '8', 'unit_count': '4'}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
response = self.client.post('/api/v1/stock-change/relaxed/', payload, format='json')
|
||||
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
|
||||
self.assertEqual(
|
||||
response.data['error'],
|
||||
f'产品ID {self.other_product.id} 对当前用户不可见',
|
||||
)
|
||||
|
||||
def test_create_stock_change_relaxed_no_employee_permission_denied(self):
|
||||
class DummyUser:
|
||||
def __init__(self, username):
|
||||
self.username = username
|
||||
self.is_authenticated = True
|
||||
|
||||
dummy_user = DummyUser('relaxed_no_emp')
|
||||
self.client.force_authenticate(user=dummy_user)
|
||||
payload = {
|
||||
'type': stock_models.StockChangeTypeEnum.ADD,
|
||||
'warehouse': self.relaxed_warehouse.id,
|
||||
'source_type': stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
'products': [
|
||||
{
|
||||
'product': self.product.id,
|
||||
'quantity': {'value': '10', 'unit_count': '5'}
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
logging.disable(logging.NOTSET)
|
||||
self.addCleanup(logging.disable, logging.CRITICAL)
|
||||
with self.assertLogs('api_v1.views.stock_change_views.mixins', level='ERROR') as cm:
|
||||
response = self.client.post('/api/v1/stock-change/relaxed/', payload, format='json')
|
||||
|
||||
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
|
||||
self.assertEqual(response.data['error'], '无权限访问')
|
||||
self.assertTrue(any('无员工信息' in msg for msg in cm.output))
|
||||
|
||||
|
||||
class StockChangeRestrictAPITestCase(TestCase):
|
||||
"""测试严进严出模式的 API"""
|
||||
|
||||
@@ -132,7 +406,7 @@ class ListStockChangeDetailsAPITestCase(TestCase):
|
||||
self.warehouse = basic_models.WareHouse.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='测试仓库',
|
||||
mode=basic_models.WareHouseModeEnum.UNRESTRICTED,
|
||||
mode=basic_models.WareHouseModeEnum.RESTRICT_IN,
|
||||
)
|
||||
|
||||
# 创建另一个产品和仓库用于测试过滤
|
||||
@@ -146,7 +420,7 @@ class ListStockChangeDetailsAPITestCase(TestCase):
|
||||
self.warehouse2 = basic_models.WareHouse.objects.create(
|
||||
merchant=self.merchant,
|
||||
name='测试仓库2',
|
||||
mode=basic_models.WareHouseModeEnum.UNRESTRICTED,
|
||||
mode=basic_models.WareHouseModeEnum.RESTRICT_IN,
|
||||
)
|
||||
|
||||
self.user = User.objects.create_user(username='detail_user', password='pass123')
|
||||
@@ -209,9 +483,9 @@ class ListStockChangeDetailsAPITestCase(TestCase):
|
||||
)
|
||||
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
||||
|
||||
# 应该返回仓库1中产品1的所有明细（record1和record2的明细）
|
||||
self.assertEqual(response.data['count'], 3) # record1有2条明细，record2有1条明细
|
||||
self.assertEqual(len(response.data['results']), 3)
|
||||
# 默认 direction=inbound，因此仅包含仓库1中产品1的入库明细（record1 的 2 条）
|
||||
self.assertEqual(response.data['count'], 2)
|
||||
self.assertEqual(len(response.data['results']), 2)
|
||||
|
||||
# 检查返回数据结构
|
||||
result = response.data['results'][0]
|
||||
|
||||
86
business/services.py
Normal file
86
business/services.py
Normal file
@@ -0,0 +1,86 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from django.db import transaction
|
||||
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from django.db import transaction
|
||||
|
||||
from basic_info import models as basic_info_models
|
||||
|
||||
from . import models
|
||||
from .tasks import create_purchase_order_stock_entries
|
||||
|
||||
|
||||
def _normalize_order_date(value) -> date:
|
||||
if isinstance(value, date):
|
||||
return value
|
||||
if isinstance(value, datetime):
|
||||
return value.date()
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return date.fromisoformat(value)
|
||||
except ValueError as exc:
|
||||
raise ValueError('order_date 格式不正确，应为 YYYY-MM-DD') from exc
|
||||
raise ValueError('order_date 格式不正确')
|
||||
|
||||
|
||||
def create_purchase_order(
|
||||
*,
|
||||
merchant: basic_info_models.Merchant,
|
||||
supplier: basic_info_models.Supplier,
|
||||
order_date,
|
||||
total_amount,
|
||||
warehouse_id: int,
|
||||
items: List[Dict[str, Any]],
|
||||
remarks: str | None = '',
|
||||
created_by=None,
|
||||
) -> models.PurchaseOrder:
|
||||
"""
|
||||
创建采购订单并触发异步创建入库单任务。
|
||||
|
||||
Args:
|
||||
merchant: 采购单所属商户
|
||||
supplier: 供应商
|
||||
order_date: 订单日期 (date)
|
||||
total_amount: 总金额
|
||||
warehouse_id: 入库仓库ID
|
||||
items: 产品明细，格式示例：
|
||||
[
|
||||
{'product_id': 1, 'quantities': ['10.00', '5.00']},
|
||||
{'product_id': 2, 'quantities': ['3.50']},
|
||||
]
|
||||
该结构会被传递给 StockFlowService，需满足其模式要求。
|
||||
remarks: 备注
|
||||
created_by: 创建者用户（可选，用于 stock 记录中的 created_by）
|
||||
"""
|
||||
if not items:
|
||||
raise ValueError('items 不能为空')
|
||||
if not warehouse_id:
|
||||
raise ValueError('warehouse_id 不能为空')
|
||||
|
||||
normalized_date = _normalize_order_date(order_date)
|
||||
total_amount = Decimal(str(total_amount))
|
||||
|
||||
with transaction.atomic():
|
||||
purchase_order = models.PurchaseOrder.objects.create(
|
||||
merchant=merchant,
|
||||
supplier=supplier,
|
||||
order_date=normalized_date,
|
||||
total_amount=total_amount,
|
||||
remarks=remarks,
|
||||
)
|
||||
|
||||
created_by_id = getattr(created_by, 'id', None)
|
||||
create_purchase_order_stock_entries.delay(
|
||||
purchase_order_id=purchase_order.id,
|
||||
warehouse_id=warehouse_id,
|
||||
items=items,
|
||||
created_by_id=created_by_id,
|
||||
)
|
||||
return purchase_order
|
||||
|
||||
56
business/tasks.py
Normal file
56
business/tasks.py
Normal file
@@ -0,0 +1,56 @@
|
||||
import logging
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from celery import shared_task
|
||||
from django.contrib.auth import get_user_model
|
||||
|
||||
from business import models as business_models
|
||||
from stock import models as stock_models
|
||||
from stock.services import StockFlowService
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@shared_task(bind=True)
|
||||
def create_purchase_order_stock_entries(
|
||||
self,
|
||||
*,
|
||||
purchase_order_id: int,
|
||||
warehouse_id: int,
|
||||
items: List[Dict[str, Any]],
|
||||
created_by_id: int | None = None,
|
||||
):
|
||||
"""
|
||||
为采购单创建入库记录。仅支持入库（StockFlowService.stock_in）。
|
||||
"""
|
||||
try:
|
||||
purchase_order = business_models.PurchaseOrder.objects.select_related('merchant').get(id=purchase_order_id)
|
||||
except business_models.PurchaseOrder.DoesNotExist:
|
||||
logger.error('PurchaseOrder %s 不存在，无法创建入库单', purchase_order_id)
|
||||
return {'error': 'purchase_order_not_found', 'purchase_order_id': purchase_order_id}
|
||||
|
||||
merchant = purchase_order.merchant
|
||||
|
||||
created_by = None
|
||||
if created_by_id:
|
||||
UserModel = get_user_model()
|
||||
created_by = UserModel.objects.filter(id=created_by_id).first()
|
||||
|
||||
service = StockFlowService(merchant=merchant, created_by=created_by)
|
||||
record, details, created_count = service.stock_in(
|
||||
warehouse_id=warehouse_id,
|
||||
source_type=stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
source_id=purchase_order.id,
|
||||
items=items,
|
||||
)
|
||||
|
||||
payload = {
|
||||
'task_id': self.request.id,
|
||||
'purchase_order_id': purchase_order.id,
|
||||
'stock_change_record_id': getattr(record, 'id', None),
|
||||
'created_details_count': created_count,
|
||||
}
|
||||
logger.info('采购单 %s 入库任务完成: %s', purchase_order.id, payload)
|
||||
return payload
|
||||
|
||||
@@ -1,3 +1,126 @@
|
||||
from decimal import Decimal
|
||||
from typing import List, Dict, Any
|
||||
|
||||
from django.contrib.auth import get_user_model
|
||||
from django.test import TestCase, override_settings
|
||||
from django.utils import timezone
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
from basic_info import models as basic_models
|
||||
from stock import models as stock_models
|
||||
from business import models as business_models, services, tasks
|
||||
|
||||
|
||||
def create_basic_fixtures():
|
||||
merchant = basic_models.Merchant.objects.create(
|
||||
name='测试商户',
|
||||
type=basic_models.MerchantTypeEnum.FACTORY,
|
||||
)
|
||||
supplier = basic_models.Supplier.objects.create(
|
||||
merchant=merchant,
|
||||
name='测试供应商',
|
||||
)
|
||||
warehouse = basic_models.WareHouse.objects.create(
|
||||
merchant=merchant,
|
||||
name='主仓',
|
||||
mode=basic_models.WareHouseModeEnum.RESTRICT_IN,
|
||||
)
|
||||
category = basic_models.ProductCategory.objects.create(
|
||||
merchant=merchant,
|
||||
name='面料',
|
||||
product_prefix='FAB',
|
||||
)
|
||||
product = basic_models.Product.objects.create(
|
||||
merchant=merchant,
|
||||
category=category,
|
||||
name='测试面料',
|
||||
human_id='FAB-001',
|
||||
unit=basic_models.ProductUnitEnum.METER,
|
||||
)
|
||||
return merchant, supplier, warehouse, product
|
||||
|
||||
|
||||
class PurchaseOrderServiceTestCase(TestCase):
|
||||
def setUp(self):
|
||||
self.merchant, self.supplier, self.warehouse, self.product = create_basic_fixtures()
|
||||
User = get_user_model()
|
||||
self.user = User.objects.create_user(username='creator', password='pass123')
|
||||
self.items: List[Dict[str, Any]] = [
|
||||
{'product_id': self.product.id, 'quantities': ['10.00', '5.00']},
|
||||
]
|
||||
|
||||
def test_create_purchase_order_triggers_task(self):
|
||||
with patch('business.services.create_purchase_order_stock_entries.delay') as mock_delay:
|
||||
purchase_order = services.create_purchase_order(
|
||||
merchant=self.merchant,
|
||||
supplier=self.supplier,
|
||||
order_date=timezone.now().date(),
|
||||
total_amount=Decimal('100.00'),
|
||||
warehouse_id=self.warehouse.id,
|
||||
items=self.items,
|
||||
remarks='自动化测试',
|
||||
created_by=self.user,
|
||||
)
|
||||
|
||||
self.assertIsInstance(purchase_order, business_models.PurchaseOrder)
|
||||
mock_delay.assert_called_once_with(
|
||||
purchase_order_id=purchase_order.id,
|
||||
warehouse_id=self.warehouse.id,
|
||||
items=self.items,
|
||||
created_by_id=self.user.id,
|
||||
)
|
||||
|
||||
def test_create_purchase_order_without_items_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
services.create_purchase_order(
|
||||
merchant=self.merchant,
|
||||
supplier=self.supplier,
|
||||
order_date=timezone.now().date(),
|
||||
total_amount=Decimal('50.00'),
|
||||
warehouse_id=self.warehouse.id,
|
||||
items=[],
|
||||
)
|
||||
|
||||
|
||||
@override_settings(
|
||||
CELERY_TASK_ALWAYS_EAGER=True,
|
||||
CELERY_TASK_EAGER_PROPAGATES=True,
|
||||
)
|
||||
class PurchaseOrderStockTaskTestCase(TestCase):
|
||||
def setUp(self):
|
||||
self.merchant, self.supplier, self.warehouse, self.product = create_basic_fixtures()
|
||||
User = get_user_model()
|
||||
self.user = User.objects.create_user(username='task-user', password='pass123')
|
||||
self.purchase_order = business_models.PurchaseOrder.objects.create(
|
||||
merchant=self.merchant,
|
||||
supplier=self.supplier,
|
||||
order_date=timezone.now().date(),
|
||||
total_amount=Decimal('100.00'),
|
||||
)
|
||||
self.items = [{'product_id': self.product.id, 'quantities': ['8.00']}]
|
||||
|
||||
def test_task_calls_stock_flow_service(self):
|
||||
with patch('business.tasks.StockFlowService') as mock_flow_cls:
|
||||
mock_instance = mock_flow_cls.return_value
|
||||
mock_instance.stock_in.return_value = (MagicMock(id=321), [], 2)
|
||||
|
||||
async_result = tasks.create_purchase_order_stock_entries.delay(
|
||||
purchase_order_id=self.purchase_order.id,
|
||||
warehouse_id=self.warehouse.id,
|
||||
items=self.items,
|
||||
created_by_id=self.user.id,
|
||||
)
|
||||
payload = async_result.get(timeout=5)
|
||||
|
||||
mock_flow_cls.assert_called_once_with(merchant=self.merchant, created_by=self.user)
|
||||
mock_instance.stock_in.assert_called_once_with(
|
||||
warehouse_id=self.warehouse.id,
|
||||
source_type=stock_models.StockChangeSourceEnum.PURCHASE,
|
||||
source_id=self.purchase_order.id,
|
||||
items=self.items,
|
||||
)
|
||||
self.assertEqual(payload['purchase_order_id'], self.purchase_order.id)
|
||||
self.assertEqual(payload['stock_change_record_id'], 321)
|
||||
from django.test import TestCase
|
||||
|
||||
# Create your tests here.
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
from django.shortcuts import render
|
||||
|
||||
# Create your views here.
|
||||
4525
data-bak/db-backup-20251126-124636.sql
Normal file
4525
data-bak/db-backup-20251126-124636.sql
Normal file
File diff suppressed because it is too large
Load Diff
76
docs/business_purchase.md
Normal file
76
docs/business_purchase.md
Normal file
@@ -0,0 +1,76 @@
|
||||
# Business 模块采购单 API
|
||||
|
||||
## 创建采购单并触发入库
|
||||
|
||||
- **URL**: `POST /api/v1/purchase-orders/`
|
||||
- **权限**: 需要登录且具备员工身份
|
||||
- **描述**: 创建业务模块的 `PurchaseOrder`，并异步触发库存入库任务（通过 `stock.services.StockFlowService.stock_in` 实现）。
|
||||
|
||||
### 请求体
|
||||
|
||||
```json
|
||||
{
|
||||
"supplier": 1,
|
||||
"warehouse": 2,
|
||||
"order_date": "2025-11-26",
|
||||
"total_amount": "1200.50",
|
||||
"remarks": "测试采购单",
|
||||
"items": [
|
||||
{
|
||||
"product_id": 10,
|
||||
"quantities": ["10.5", "6.0"]
|
||||
},
|
||||
{
|
||||
"product_id": 18,
|
||||
"quantities": ["3.25"]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
| 字段 | 类型 | 必填 | 说明 |
|
||||
|------|------|------|------|
|
||||
| supplier | integer | ✅ | 供应商 ID，必须隶属于当前商户 |
|
||||
| warehouse | integer | ✅ | 入库仓库 ID |
|
||||
| order_date | string (date) | ✅ | 订单日期（`YYYY-MM-DD`） |
|
||||
| total_amount | string/number | ✅ | 采购总金额 |
|
||||
| remarks | string | 否 | 备注 |
|
||||
| items | array | ✅ | 入库明细，结构需满足 `StockFlowService` 的要求（目前仅支持严谨模式：product_id + quantities） |
|
||||
|
||||
> `items` 内部字段会被直接传给 `StockFlowService`，因此：
|
||||
> - 严谨/严进模式：`{'product_id': 1, 'quantities': ['10', '5']}`
|
||||
> - 宽松模式/严进严出出库暂未开放，后续按需扩展。
|
||||
|
||||
### 响应
|
||||
|
||||
```json
|
||||
{
|
||||
"id": 35,
|
||||
"message": "采购单创建成功，入库任务已排队"
|
||||
}
|
||||
```
|
||||
|
||||
创建成功即刻返回，实际入库明细由后台 Celery 任务生成，可在日志或 `stock_change` 记录中查看。
|
||||
|
||||
### 错误示例
|
||||
|
||||
| 状态码 | 示例 | 说明 |
|
||||
|--------|------|------|
|
||||
| 400 | `{"error": "缺少供应商 ID"}` | 请求缺失关键字段 |
|
||||
| 400 | `{"error": "供应商 99 不存在"}` | 提供的供应商不属于当前商户 |
|
||||
| 403 | `{"error": "无权限访问"}` | 当前用户无员工信息 |
|
||||
|
||||
### 关联任务（business/tasks.py）
|
||||
|
||||
`create_purchase_order_stock_entries` 任务会接收 `purchase_order_id`、`warehouse_id`、`items` 等信息，并通过 `StockFlowService.stock_in` 创建入库记录。
|
||||
|
||||
- 任务日志示例：`采购单 35 入库任务完成`
|
||||
- 返回 payload 包含 `stock_change_record_id`、`created_details_count`
|
||||
|
||||
### 测试
|
||||
|
||||
`api_v1/tests.py` 中新增 `PurchaseOrderAPITestCase`，通过 eager Celery 设置验证：
|
||||
|
||||
1. API 请求返回 201
|
||||
2. Celery 任务被成功调用（patch `create_purchase_order_stock_entries.delay` 断言参数）
|
||||
|
||||
108
docs/celery_testing.md
Normal file
108
docs/celery_testing.md
Normal file
@@ -0,0 +1,108 @@
|
||||
# Celery 任务测试指南
|
||||
|
||||
## 1. 自动化测试（单元测试/CI）
|
||||
|
||||
项目新增了三个演示型任务（`api_v1/tasks.py`）：
|
||||
|
||||
| 任务 | 功能 |
|
||||
|------|------|
|
||||
| `ping_task(message='ping')` | 打印并返回时间戳，用于快速验证 worker 是否可接收/返回结果 |
|
||||
| `merchant_product_count(merchant_id)` | 统计某商户下产品数量，示范如何在任务中访问数据库 |
|
||||
| `backup_database(output_dir=None, filename_prefix='db-backup')` | 将当前数据库导出为 `.sql` 文件（数据备份），默认保存到 `BASE_DIR/data-bak/` |
|
||||
|
||||
对应的单元测试位于 `api_v1/tests.py`（类 `CeleryTasksTestCase`），通过
|
||||
`@override_settings(CELERY_TASK_ALWAYS_EAGER=True)` 让任务在测试进程内同步执行。
|
||||
|
||||
运行方式：
|
||||
|
||||
```bash
|
||||
uv run manage.py test api_v1.tests.CeleryTasksTestCase
|
||||
```
|
||||
|
||||
若要在其他测试中调用任务，只需在测试类上使用同样的 `override_settings`
|
||||
即可保证 Celery 在没有 worker 的情况下仍能同步执行。
|
||||
|
||||
---
|
||||
|
||||
## 2. 本地验证真实 worker（RabbitMQ + Celery）
|
||||
|
||||
### 2.1 启动依赖服务
|
||||
|
||||
```bash
|
||||
# 启动基础设施（Postgres/Redis/Rabbit/Celery worker/web）
|
||||
docker compose up -d postgres redis rabbitmq
|
||||
|
||||
# 启动 web（Django）
|
||||
docker compose up -d web
|
||||
|
||||
# 启动 Celery worker（如已运行，可跳过）
|
||||
docker compose up -d celery_worker
|
||||
# 或者手动：uv run celery -A flower worker -l info
|
||||
```
|
||||
|
||||
> 若 worker 容器异常，可通过 `docker compose restart celery_worker` 重启。
|
||||
|
||||
### 2.2 发送测试任务
|
||||
|
||||
在另一个终端进入容器或宿主项目目录执行：
|
||||
|
||||
```bash
|
||||
uv run python manage.py shell
|
||||
```
|
||||
|
||||
```python
|
||||
from api_v1.tasks import ping_task, merchant_product_count, backup_database
|
||||
from basic_info.models import Merchant
|
||||
|
||||
# 发送心跳任务
|
||||
async_result = ping_task.delay('hello celery')
|
||||
print(async_result.get(timeout=10))
|
||||
|
||||
# 发送统计任务（示例：使用 ID=1 的商户）
|
||||
merchant = Merchant.objects.first()
|
||||
result = merchant_product_count.delay(merchant.id)
|
||||
print(result.get(timeout=10))
|
||||
|
||||
# 触发数据库备份（输出到默认 data-bak）
|
||||
backup = backup_database.delay(filename_prefix='manual-backup')
|
||||
print(backup.get(timeout=30)) # payload 中包含 backup_path，文件为 .sql
|
||||
```
|
||||
|
||||
### 2.3 观察执行结果
|
||||
|
||||
1. **Worker 日志**：`docker compose logs -f celery_worker` \
|
||||
可查看任务被消费、日志输出等信息。
|
||||
2. **RabbitMQ 控制台**：访问 `http://localhost:15672`（默认账号 guest/guest），
|
||||
观察队列长度是否回落到 0。
|
||||
3. **Task Result**：上面的 `result.get()` 会在任务完成时返回 payload，
|
||||
若超时或无法连接则会抛出异常，帮助定位问题。
|
||||
|
||||
---
|
||||
|
||||
## 3. 失败排查 Checklist
|
||||
|
||||
1. **环境变量**：`CELERY_BROKER_URL` 和 `CELERY_RESULT_BACKEND` 是否指向
|
||||
正在运行的 RabbitMQ/Redis？
|
||||
2. **Worker 进程**：`docker compose ps` 确认 `celery_worker` 状态为 Up。
|
||||
3. **队列阻塞**：RabbitMQ 控制台查看是否有大量消息处于 `Unacked`。
|
||||
4. **日志级别**：测试中需要捕获日志时，可使用
|
||||
`with self.assertLogs('api_v1.tasks', level='INFO')`.
|
||||
5. **自动化测试**：若任务在测试里需要真实队列，请去掉
|
||||
`CELERY_TASK_ALWAYS_EAGER` 覆盖；否则保持默认值即可同步执行。
|
||||
6. **数据库备份依赖**：`backup_database` 在 PostgreSQL 场景下需要系统可执行 `pg_dump`，
|
||||
请确保容器/宿主机已安装 PostgreSQL 客户端工具；SQLite 则会使用内置 `iterdump` 生成 `.sql`。
|
||||
|
||||
---
|
||||
|
||||
## 4. 常用命令速查
|
||||
|
||||
| 操作 | 命令 |
|
||||
|------|------|
|
||||
| 启动 worker | `docker compose up -d celery_worker` |
|
||||
| 查看 worker 日志 | `docker compose logs -f celery_worker` |
|
||||
| 重启 worker | `docker compose restart celery_worker` |
|
||||
| 清空 RabbitMQ 队列 | `docker compose exec rabbitmq rabbitmqctl purge_queue flower`（示例） |
|
||||
| 运行单个任务测试 | `uv run python manage.py shell` -> `ping_task.delay()` |
|
||||
|
||||
通过以上步骤，即可确认 Celery + RabbitMQ 在本地能够成功执行任务，并在自动化测试里保持覆盖。若需要新增业务任务，可参考 `api_v1/tasks.py` 的写法：使用 `@shared_task`，在任务内进行必要的日志记录，便于问题排查。***
|
||||
|
||||
217
docs/sse.md
Normal file
217
docs/sse.md
Normal file
@@ -0,0 +1,217 @@
|
||||
# SSE (Server-Sent Events) 模块
|
||||
|
||||
基于 Django 和 Django REST Framework 的服务器推送事件实现，支持多商户隔离。
|
||||
|
||||
## 功能特性
|
||||
|
||||
- ✅ 使用 DRF 处理请求和响应（非流式端点）
|
||||
- ✅ 支持多种请求格式（JSON、Form Data、Multipart）
|
||||
- ✅ 自动数据验证和序列化
|
||||
- ✅ 异步支持，高并发处理
|
||||
- ✅ 心跳机制保持连接活跃
|
||||
- ✅ 自动清理断开的连接
|
||||
- ✅ 连接状态监控
|
||||
- ✅ 多商户消息隔离
|
||||
- ✅ JWT认证和商户验证
|
||||
|
||||
## API 端点
|
||||
|
||||
### 1. 订阅 SSE 事件流
|
||||
|
||||
**端点**: `GET /sse/`
|
||||
|
||||
客户端连接此端点保持长连接，接收服务器推送的事件。
|
||||
|
||||
**认证要求**: 需要JWT认证，且用户必须有关联的商户
|
||||
|
||||
**示例**:
|
||||
```bash
|
||||
# 需要提供JWT token
|
||||
curl -N -H "Authorization: Bearer YOUR_JWT_TOKEN" http://localhost:8000/sse/
|
||||
```
|
||||
|
||||
**JavaScript 示例**:
|
||||
```javascript
|
||||
// 需要在连接时提供认证头
|
||||
const eventSource = new EventSource('/sse/', {
|
||||
headers: {
|
||||
'Authorization': `Bearer ${yourJwtToken}`
|
||||
}
|
||||
});
|
||||
|
||||
eventSource.onmessage = function(event) {
|
||||
const data = JSON.parse(event.data);
|
||||
console.log('收到消息:', data);
|
||||
console.log('商户ID:', data.merchant_id);
|
||||
};
|
||||
```
|
||||
|
||||
**注意**: 浏览器原生的EventSource不支持自定义请求头，因此推荐使用EventSource polyfill或fetch实现。
|
||||
|
||||
---
|
||||
|
||||
### 2. 推送事件到当前商户的客户端
|
||||
|
||||
**端点**: `POST /sse/push/`
|
||||
|
||||
向当前用户所属商户的所有已连接客户端广播消息。
|
||||
|
||||
**认证要求**: 需要JWT认证，且用户必须有关联的商户
|
||||
|
||||
**请求参数**:
|
||||
- 自动使用当前用户的商户ID
|
||||
- 固定发送测试消息：'订单已支付'
|
||||
- 固定事件类型：'order_paid'
|
||||
- 固定对象ID：12345
|
||||
|
||||
**请求示例**:
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/sse/push/ \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer YOUR_JWT_TOKEN"
|
||||
```
|
||||
|
||||
**响应示例**:
|
||||
```json
|
||||
{
|
||||
"status": "ok",
|
||||
"message": "Test event broadcasted to your merchant",
|
||||
"merchant_id": 1,
|
||||
"clients": 3
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### 3. 获取连接状态
|
||||
|
||||
**端点**: `GET /sse/status/`
|
||||
|
||||
查询当前 SSE 服务器的状态和连接数。
|
||||
|
||||
**认证要求**: 需要JWT认证，且用户必须有关联的商户
|
||||
|
||||
**示例**:
|
||||
```bash
|
||||
curl -H "Authorization: Bearer YOUR_JWT_TOKEN" http://localhost:8000/sse/status/
|
||||
```
|
||||
|
||||
**响应示例**:
|
||||
```json
|
||||
{
|
||||
"status": "running",
|
||||
"total_clients": 10,
|
||||
"merchant_clients": 3,
|
||||
"merchant_id": 1,
|
||||
"message": "SSE server is running with 10 total connections, 3 for your merchant"
|
||||
}
|
||||
```
|
||||
|
||||
### 4. 关闭商户连接
|
||||
|
||||
**端点**: `POST /sse/shutdown/`
|
||||
|
||||
关闭当前用户所属商户的所有SSE连接。
|
||||
|
||||
**认证要求**: 需要JWT认证，且用户必须有关联的商户
|
||||
|
||||
**示例**:
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/sse/shutdown/ \
|
||||
-H "Authorization: Bearer YOUR_JWT_TOKEN"
|
||||
```
|
||||
|
||||
**响应示例**:
|
||||
```json
|
||||
{
|
||||
"status": "ok",
|
||||
"message": "Shutdown signal sent to your merchant's SSE connections",
|
||||
"merchant_id": 1,
|
||||
"clients": 3
|
||||
}
|
||||
```
|
||||
|
||||
## 启动服务器
|
||||
|
||||
使用 Uvicorn (ASGI 服务器) 启动：
|
||||
|
||||
```bash
|
||||
# 开发环境
|
||||
uvicorn flower.asgi:application --reload --host 0.0.0.0 --port 8000
|
||||
|
||||
# 生产环境
|
||||
uvicorn flower.asgi:application --host 0.0.0.0 --port 8000 --workers 4
|
||||
```
|
||||
|
||||
## 测试
|
||||
|
||||
打开 `sse_test.html` 在浏览器中测试：
|
||||
|
||||
1. 点击"连接 SSE"建立连接
|
||||
2. 输入消息
|
||||
3. 点击"发送 (JSON)"或"发送 (Form Data)"测试不同格式
|
||||
4. 点击"获取连接状态"查看当前连接数
|
||||
5. 打开多个浏览器标签测试广播功能
|
||||
|
||||
## Python 客户端示例
|
||||
|
||||
```python
|
||||
import requests
|
||||
import sseclient # pip install sseclient-py
|
||||
|
||||
# 订阅事件
|
||||
response = requests.get('http://localhost:8000/sse/', stream=True)
|
||||
client = sseclient.SSEClient(response)
|
||||
|
||||
for event in client.events():
|
||||
print(f'收到消息: {event.data}')
|
||||
```
|
||||
|
||||
## 技术实现
|
||||
|
||||
- **流式响应**: 使用 `StreamingHttpResponse` 实现 SSE 连接
|
||||
- **队列机制**: 每个连接对应一个 `queue.Queue`，按商户组织
|
||||
- **心跳**: 30 秒超时，自动发送心跳保持连接
|
||||
- **DRF 集成**: 非流式端点使用 DRF 的 `@api_view` 和序列化器
|
||||
- **多格式支持**: 自动解析 JSON、Form Data、Multipart 等格式
|
||||
- **商户隔离**: 所有消息按商户隔离，确保数据安全
|
||||
- **JWT认证**: 使用 DRF Simple JWT 进行身份验证
|
||||
|
||||
## 消息结构
|
||||
|
||||
### 连接成功消息
|
||||
```json
|
||||
{
|
||||
"type": "connected",
|
||||
"message": "SSE connection established",
|
||||
"merchant_id": 1
|
||||
}
|
||||
```
|
||||
|
||||
### 业务事件消息
|
||||
```json
|
||||
{
|
||||
"mode": "simple_message",
|
||||
"type": "order_paid",
|
||||
"message": "订单已支付",
|
||||
"object_id": 12345,
|
||||
"merchant_id": 1
|
||||
}
|
||||
```
|
||||
|
||||
### 服务器关闭消息
|
||||
```json
|
||||
{
|
||||
"type": "server_shutdown",
|
||||
"message": "Server shutting down your connections, please reconnect later"
|
||||
}
|
||||
```
|
||||
|
||||
## 注意事项
|
||||
|
||||
1. 必须使用 ASGI 服务器（如 Uvicorn、Daphne）运行
|
||||
2. 不支持使用传统的 WSGI 服务器（如 Gunicorn + WSGI）
|
||||
3. 如果使用 Nginx，需要禁用缓冲：`X-Accel-Buffering: no`
|
||||
4. 所有SSE端点都需要JWT认证，且用户必须有关联的商户
|
||||
5. 浏览器原生的EventSource不支持自定义请求头，推荐使用polyfill或fetch实现
|
||||
6. SSE 使用 GET 请求，注意 CORS 配置
|
||||
@@ -12,4 +12,3 @@ app.autodiscover_tasks()
|
||||
@app.task(bind=True)
|
||||
def debug_task(self):
|
||||
print(f'Celery debug task - request: {self.request!r}')
|
||||
|
||||
|
||||
154
sse/README.md
154
sse/README.md
@@ -1,154 +0,0 @@
|
||||
# SSE (Server-Sent Events) 模块
|
||||
|
||||
基于 Django REST Framework 的服务器推送事件实现。
|
||||
|
||||
## 功能特性
|
||||
|
||||
- ✅ 使用 DRF 处理请求和响应
|
||||
- ✅ 支持多种请求格式（JSON、Form Data、Multipart）
|
||||
- ✅ 自动数据验证和序列化
|
||||
- ✅ 异步支持，高并发处理
|
||||
- ✅ 心跳机制保持连接活跃
|
||||
- ✅ 自动清理断开的连接
|
||||
- ✅ 连接状态监控
|
||||
|
||||
## API 端点
|
||||
|
||||
### 1. 订阅 SSE 事件流
|
||||
|
||||
**端点**: `GET /sse/`
|
||||
|
||||
客户端连接此端点保持长连接，接收服务器推送的事件。
|
||||
|
||||
**示例**:
|
||||
```bash
|
||||
curl -N http://localhost:8000/sse/
|
||||
```
|
||||
|
||||
**JavaScript 示例**:
|
||||
```javascript
|
||||
const eventSource = new EventSource('http://localhost:8000/sse/');
|
||||
|
||||
eventSource.onmessage = function(event) {
|
||||
const data = JSON.parse(event.data);
|
||||
console.log('收到消息:', data);
|
||||
};
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### 2. 推送事件到所有客户端
|
||||
|
||||
**端点**: `POST /sse/push/`
|
||||
|
||||
向所有已连接的客户端广播消息。
|
||||
|
||||
**请求参数**:
|
||||
- `message` (必填): 消息内容
|
||||
- `type` (可选): 事件类型，默认为 'message'
|
||||
|
||||
**支持的请求格式**:
|
||||
|
||||
#### JSON 格式
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/sse/push/ \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"message": "Hello, SSE!", "type": "notification"}'
|
||||
```
|
||||
|
||||
#### Form Data 格式
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/sse/push/ \
|
||||
-d "message=Hello, SSE!" \
|
||||
-d "type=notification"
|
||||
```
|
||||
|
||||
#### Multipart Form Data
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/sse/push/ \
|
||||
-F "message=Hello, SSE!" \
|
||||
-F "type=notification"
|
||||
```
|
||||
|
||||
**响应示例**:
|
||||
```json
|
||||
{
|
||||
"status": "success",
|
||||
"message": "Event sent to 3 client(s)",
|
||||
"clients": 3,
|
||||
"sent": 3
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### 3. 获取连接状态
|
||||
|
||||
**端点**: `GET /sse/status/`
|
||||
|
||||
查询当前 SSE 服务器的状态和连接数。
|
||||
|
||||
**示例**:
|
||||
```bash
|
||||
curl http://localhost:8000/sse/status/
|
||||
```
|
||||
|
||||
**响应示例**:
|
||||
```json
|
||||
{
|
||||
"status": "running",
|
||||
"clients": 3,
|
||||
"message": "SSE server is running with 3 active connection(s)"
|
||||
}
|
||||
```
|
||||
|
||||
## 启动服务器
|
||||
|
||||
使用 Uvicorn (ASGI 服务器) 启动：
|
||||
|
||||
```bash
|
||||
# 开发环境
|
||||
uvicorn flower.asgi:application --reload --host 0.0.0.0 --port 8000
|
||||
|
||||
# 生产环境
|
||||
uvicorn flower.asgi:application --host 0.0.0.0 --port 8000 --workers 4
|
||||
```
|
||||
|
||||
## 测试
|
||||
|
||||
打开 `sse_test.html` 在浏览器中测试：
|
||||
|
||||
1. 点击"连接 SSE"建立连接
|
||||
2. 输入消息
|
||||
3. 点击"发送 (JSON)"或"发送 (Form Data)"测试不同格式
|
||||
4. 点击"获取连接状态"查看当前连接数
|
||||
5. 打开多个浏览器标签测试广播功能
|
||||
|
||||
## Python 客户端示例
|
||||
|
||||
```python
|
||||
import requests
|
||||
import sseclient # pip install sseclient-py
|
||||
|
||||
# 订阅事件
|
||||
response = requests.get('http://localhost:8000/sse/', stream=True)
|
||||
client = sseclient.SSEClient(response)
|
||||
|
||||
for event in client.events():
|
||||
print(f'收到消息: {event.data}')
|
||||
```
|
||||
|
||||
## 技术实现
|
||||
|
||||
- **异步视图**: 使用 `async def` 实现异步处理
|
||||
- **队列机制**: 每个连接对应一个 `asyncio.Queue`
|
||||
- **心跳**: 30 秒超时，自动发送心跳保持连接
|
||||
- **DRF 集成**: 使用 DRF 的 `@api_view` 和序列化器
|
||||
- **多格式支持**: 自动解析 JSON、Form Data、Multipart 等格式
|
||||
|
||||
## 注意事项
|
||||
|
||||
1. 必须使用 ASGI 服务器（如 Uvicorn、Daphne）运行
|
||||
2. 不支持使用传统的 WSGI 服务器（如 Gunicorn + WSGI）
|
||||
3. 如果使用 Nginx，需要禁用缓冲：`X-Accel-Buffering: no`
|
||||
4. SSE 使用 GET 请求，注意 CORS 配置
|
||||
117
sse/auth_utils.py
Normal file
117
sse/auth_utils.py
Normal file
@@ -0,0 +1,117 @@
|
||||
"""
|
||||
SSE认证工具模块
|
||||
|
||||
提供统一的认证和商户验证功能，用于SSE连接端点和其他需要认证的地方
|
||||
"""
|
||||
|
||||
from django.http import HttpResponseForbidden
|
||||
from rest_framework.authentication import get_authorization_header
|
||||
from rest_framework_simplejwt.authentication import JWTAuthentication
|
||||
from rest_framework_simplejwt.exceptions import InvalidToken
|
||||
from rest_framework.exceptions import AuthenticationFailed
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_user_merchant_id(request):
|
||||
"""获取用户所属商户ID"""
|
||||
try:
|
||||
# 检查用户是否已认证
|
||||
if not request.user.is_authenticated:
|
||||
return None
|
||||
# 获取商户ID
|
||||
return request.user.employee.merchant_id
|
||||
except (AttributeError, AttributeError):
|
||||
return None
|
||||
|
||||
|
||||
def authenticate_sse_request(request):
|
||||
"""
|
||||
对SSE连接请求进行认证
|
||||
|
||||
参数:
|
||||
- request: Django HTTP请求对象
|
||||
|
||||
返回:
|
||||
- tuple: (user, merchant_id) 认证成功返回用户和商户ID
|
||||
- None: 认证失败返回None
|
||||
"""
|
||||
jwt_authenticator = JWTAuthentication()
|
||||
|
||||
try:
|
||||
# 获取Authorization头
|
||||
auth_header = get_authorization_header(request).split()
|
||||
if not auth_header or auth_header[0].lower() != b'bearer':
|
||||
logger.warning("SSE认证失败: 缺少Bearer token")
|
||||
return None
|
||||
|
||||
if len(auth_header) < 2:
|
||||
logger.warning("SSE认证失败: Bearer token格式错误")
|
||||
return None
|
||||
|
||||
token = auth_header[1].decode('utf-8')
|
||||
|
||||
# 验证token
|
||||
validated_token = jwt_authenticator.get_validated_token(token)
|
||||
user = jwt_authenticator.get_user(validated_token)
|
||||
|
||||
# 检查用户是否有关联的员工和商户
|
||||
try:
|
||||
merchant_id = user.employee.merchant_id
|
||||
if not merchant_id:
|
||||
logger.warning(f"用户 {user.username} 无关联商户")
|
||||
return None
|
||||
return user, merchant_id
|
||||
except AttributeError:
|
||||
logger.warning(f"用户 {user.username} 无关联员工")
|
||||
return None
|
||||
|
||||
except (IndexError, InvalidToken, AuthenticationFailed) as e:
|
||||
logger.warning(f"SSE认证失败: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def require_sse_authentication(view_func):
|
||||
"""
|
||||
装饰器：要求SSE连接认证
|
||||
|
||||
用于SSE连接端点的认证，如果认证失败则返回403响应
|
||||
OPTIONS请求跳过认证
|
||||
"""
|
||||
def wrapper(request, *args, **kwargs):
|
||||
# OPTIONS请求跳过认证
|
||||
if request.method == 'OPTIONS':
|
||||
return view_func(request, *args, **kwargs)
|
||||
|
||||
auth_result = authenticate_sse_request(request)
|
||||
if not auth_result:
|
||||
return HttpResponseForbidden("Authentication failed or user has no associated merchant")
|
||||
|
||||
user, merchant_id = auth_result
|
||||
request.user = user
|
||||
request.merchant_id = merchant_id
|
||||
|
||||
return view_func(request, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def validate_sse_request(request):
|
||||
"""
|
||||
验证SSE请求的认证状态
|
||||
|
||||
返回:
|
||||
- tuple: (is_valid, error_response)
|
||||
- is_valid: bool 表示请求是否有效
|
||||
- error_response: HttpResponse 如果无效则返回错误响应，否则为None
|
||||
"""
|
||||
auth_result = authenticate_sse_request(request)
|
||||
if not auth_result:
|
||||
return False, HttpResponseForbidden("认证失败或用户无关联商户")
|
||||
|
||||
user, merchant_id = auth_result
|
||||
request.user = user
|
||||
request.merchant_id = merchant_id
|
||||
|
||||
return True, None
|
||||
187
sse/client_example.js
Normal file
187
sse/client_example.js
Normal file
@@ -0,0 +1,187 @@
|
||||
/**
|
||||
* SSE客户端示例
|
||||
*
|
||||
* 由于浏览器原生的EventSource不支持自定义请求头，
|
||||
* 这里提供了两种实现方式：
|
||||
* 1. 使用fetch实现（推荐）
|
||||
* 2. 使用EventSource polyfill（备选）
|
||||
*/
|
||||
|
||||
/**
|
||||
* 方式1：使用fetch实现SSE客户端（推荐）
|
||||
*
|
||||
* @param {string} url - SSE端点URL
|
||||
* @param {string} token - JWT认证token
|
||||
* @param {function} onMessage - 接收到消息时的回调函数
|
||||
* @param {function} onError - 连接错误时的回调函数
|
||||
* @param {function} onClose - 连接关闭时的回调函数
|
||||
*/
|
||||
function createSSEConnectionWithFetch(url, token, onMessage, onError, onClose) {
|
||||
const controller = new AbortController();
|
||||
const signal = controller.signal;
|
||||
|
||||
// 启动一个长时间运行的fetch请求
|
||||
fetch(url, {
|
||||
method: 'GET',
|
||||
headers: {
|
||||
'Authorization': `Bearer ${token}`,
|
||||
'Accept': 'text/event-stream',
|
||||
'Cache-Control': 'no-cache',
|
||||
},
|
||||
signal: signal
|
||||
})
|
||||
.then(response => {
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP error! status: ${response.status}`);
|
||||
}
|
||||
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
|
||||
function processBuffer() {
|
||||
const lines = buffer.split('\n');
|
||||
buffer = lines.pop(); // 保留最后一个不完整的行
|
||||
|
||||
for (const line of lines) {
|
||||
if (line.trim() === '') continue; // 空行表示事件结束
|
||||
|
||||
if (line.startsWith('data: ')) {
|
||||
const data = line.substring(6); // 去掉 'data: ' 前缀
|
||||
try {
|
||||
const event = JSON.parse(data);
|
||||
if (onMessage) onMessage(event);
|
||||
} catch (e) {
|
||||
console.error('Error parsing SSE data:', e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function read() {
|
||||
return reader.read().then(({ done, value }) => {
|
||||
if (done) {
|
||||
if (onClose) onClose();
|
||||
return;
|
||||
}
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
processBuffer();
|
||||
|
||||
// 继续读取
|
||||
return read();
|
||||
});
|
||||
}
|
||||
|
||||
return read();
|
||||
})
|
||||
.catch(error => {
|
||||
if (onError) onError(error);
|
||||
});
|
||||
|
||||
// 返回一个对象，包含关闭连接的方法
|
||||
return {
|
||||
close: () => controller.abort()
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* 方式2：使用EventSource Polyfill实现
|
||||
*
|
||||
* 需要先安装EventSource polyfill：
|
||||
* npm install event-source-polyfill
|
||||
*
|
||||
* 在应用入口处导入：
|
||||
* import 'event-source-polyfill';
|
||||
*/
|
||||
function createSSEConnectionWithPolyfill(url, token, onMessage, onError, onClose) {
|
||||
// 创建带有认证的URL
|
||||
const urlWithAuth = `${url}?token=${encodeURIComponent(token)}`;
|
||||
|
||||
const eventSource = new EventSource(urlWithAuth);
|
||||
|
||||
eventSource.onmessage = function(event) {
|
||||
try {
|
||||
const data = JSON.parse(event.data);
|
||||
if (onMessage) onMessage(data);
|
||||
} catch (e) {
|
||||
console.error('Error parsing SSE data:', e);
|
||||
}
|
||||
};
|
||||
|
||||
eventSource.onerror = function(error) {
|
||||
if (onError) onError(error);
|
||||
};
|
||||
|
||||
eventSource.onclose = function() {
|
||||
if (onClose) onClose();
|
||||
};
|
||||
|
||||
return {
|
||||
close: () => eventSource.close()
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* 使用示例
|
||||
*/
|
||||
const JWT_TOKEN = 'your_jwt_token_here';
|
||||
const SSE_URL = '/sse/';
|
||||
|
||||
// 使用fetch方式（推荐）
|
||||
const sseConnection = createSSEConnectionWithFetch(
|
||||
SSE_URL,
|
||||
JWT_TOKEN,
|
||||
(event) => {
|
||||
console.log('收到SSE事件:', event);
|
||||
|
||||
// 根据事件类型处理不同业务逻辑
|
||||
switch (event.type) {
|
||||
case 'connected':
|
||||
console.log(`SSE连接成功，商户ID: ${event.merchant_id}`);
|
||||
break;
|
||||
|
||||
case 'order_paid':
|
||||
console.log(`订单已支付: ${event.object_id}`);
|
||||
// 刷新订单列表或显示通知
|
||||
break;
|
||||
|
||||
case 'stock_change_record':
|
||||
console.log(`库存变动记录: ${event.object_id}`);
|
||||
// 刷新库存数据
|
||||
break;
|
||||
|
||||
case 'server_shutdown':
|
||||
console.log('服务器即将关闭连接');
|
||||
// 可以提示用户重新连接
|
||||
break;
|
||||
}
|
||||
},
|
||||
(error) => {
|
||||
console.error('SSE连接错误:', error);
|
||||
// 可以在这里实现重连逻辑
|
||||
setTimeout(() => {
|
||||
console.log('尝试重新连接...');
|
||||
// 重新创建连接
|
||||
}, 5000);
|
||||
},
|
||||
() => {
|
||||
console.log('SSE连接已关闭');
|
||||
}
|
||||
);
|
||||
|
||||
// 当需要关闭连接时（例如用户退出登录）
|
||||
// sseConnection.close();
|
||||
|
||||
/**
|
||||
* 如果需要使用polyfill方式，确保在应用入口处导入polyfill
|
||||
*/
|
||||
// import 'event-source-polyfill';
|
||||
//
|
||||
// const sseConnection = createSSEConnectionWithPolyfill(
|
||||
// SSE_URL,
|
||||
// JWT_TOKEN,
|
||||
// (event) => { /* 处理事件 */ },
|
||||
// (error) => { /* 处理错误 */ },
|
||||
// () => { /* 处理关闭 */ }
|
||||
// );
|
||||
143
sse/services.py
143
sse/services.py
@@ -3,50 +3,81 @@ import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 存储所有活动的 SSE 连接队列（同步队列）
|
||||
_connections = set()
|
||||
# 存储所有活动的 SSE 连接队列（按商户ID组织）
|
||||
_connections = {} # {merchant_id: set([conn_queue1, conn_queue2, ...])}
|
||||
|
||||
|
||||
def get_active_connections():
|
||||
"""获取当前所有活动的 SSE 连接队列"""
|
||||
return _connections
|
||||
def get_all_connections_count():
|
||||
"""获取所有商户的连接总数"""
|
||||
total = 0
|
||||
for merchant_id, connections in _connections.items():
|
||||
total += len(connections)
|
||||
return total
|
||||
|
||||
|
||||
def push_connection(conn_queue):
|
||||
def get_merchant_connections(merchant_id):
|
||||
"""获取指定商户的所有连接队列"""
|
||||
return _connections.get(merchant_id, set())
|
||||
|
||||
|
||||
def get_all_merchant_ids():
|
||||
"""获取所有有活跃连接的商户ID列表"""
|
||||
return list(_connections.keys())
|
||||
|
||||
|
||||
def push_connection(merchant_id, conn_queue):
|
||||
"""
|
||||
将一个新的连接队列添加到活动连接集合中
|
||||
将一个新的连接队列添加到指定商户的连接集合中
|
||||
|
||||
参数:
|
||||
- merchant_id: 商户ID
|
||||
- conn_queue: queue.Queue 实例
|
||||
"""
|
||||
_connections.add(conn_queue)
|
||||
logger.info(f"新的 SSE 连接建立，当前连接数: {len(_connections)}")
|
||||
if merchant_id not in _connections:
|
||||
_connections[merchant_id] = set()
|
||||
|
||||
_connections[merchant_id].add(conn_queue)
|
||||
total_connections = sum(len(conns) for conns in _connections.values())
|
||||
logger.info(f"商户 {merchant_id} 的新 SSE 连接建立，该商户连接数: {len(_connections[merchant_id])}，总连接数: {total_connections}")
|
||||
|
||||
|
||||
def remove_connection(conn_queue):
|
||||
def remove_connection(merchant_id, conn_queue):
|
||||
"""
|
||||
从活动连接集合中移除一个连接队列
|
||||
从指定商户的连接集合中移除一个连接队列
|
||||
|
||||
参数:
|
||||
- merchant_id: 商户ID
|
||||
- conn_queue: queue.Queue 实例
|
||||
"""
|
||||
_connections.discard(conn_queue)
|
||||
logger.info(f"SSE 连接断开，当前连接数: {len(_connections)}")
|
||||
if merchant_id in _connections:
|
||||
_connections[merchant_id].discard(conn_queue)
|
||||
|
||||
# 如果该商户没有其他连接，则移除整个商户记录
|
||||
if not _connections[merchant_id]:
|
||||
del _connections[merchant_id]
|
||||
|
||||
total_connections = sum(len(conns) for conns in _connections.values())
|
||||
logger.info(f"商户 {merchant_id} 的 SSE 连接断开，该商户连接数: {len(_connections.get(merchant_id, set()))}，总连接数: {total_connections}")
|
||||
else:
|
||||
logger.warning(f"尝试从未注册的商户 {merchant_id} 移除连接")
|
||||
|
||||
|
||||
def cleanup_all_connections():
|
||||
"""
|
||||
清理所有连接（用于服务器关闭时）
|
||||
"""
|
||||
connections_list = list(_connections)
|
||||
for conn_queue in connections_list:
|
||||
try:
|
||||
# 发送关闭信号到队列
|
||||
if hasattr(conn_queue, 'put_nowait'):
|
||||
conn_queue.put_nowait({'type': 'server_shutdown', 'message': 'Server shutting down'})
|
||||
logger.debug(f'shutdown connection {conn_queue}')
|
||||
except Exception:
|
||||
pass # 忽略错误，因为连接可能已经断开
|
||||
all_connections = {}
|
||||
for merchant_id, connections in list(_connections.items()):
|
||||
all_connections[merchant_id] = list(connections)
|
||||
for conn_queue in connections:
|
||||
try:
|
||||
# 发送关闭信号到队列
|
||||
if hasattr(conn_queue, 'put_nowait'):
|
||||
conn_queue.put_nowait({'type': 'server_shutdown', 'message': 'Server shutting down'})
|
||||
logger.debug(f'shutdown connection {conn_queue} for merchant {merchant_id}')
|
||||
except Exception:
|
||||
pass # 忽略错误，因为连接可能已经断开
|
||||
|
||||
_connections.clear()
|
||||
logger.info("所有SSE连接已清理")
|
||||
|
||||
@@ -58,23 +89,60 @@ def push_sse_event_to_all(event_data: dict):
|
||||
参数:
|
||||
- event_data: 要发送的事件数据（字典）
|
||||
"""
|
||||
disconnected = {}
|
||||
|
||||
for merchant_id, connections in list(_connections.items()):
|
||||
disconnected[merchant_id] = []
|
||||
|
||||
for conn_queue in connections:
|
||||
try:
|
||||
# 使用put_nowait非阻塞发送消息
|
||||
conn_queue.put_nowait(event_data)
|
||||
except queue.Full:
|
||||
# 队列满了，说明客户端处理消息太慢，跳过这条消息
|
||||
logger.warning(f"商户 {merchant_id} 队列已满，跳过消息: {event_data.get('type', 'unknown')}")
|
||||
except Exception as e:
|
||||
# 其他异常可能表示连接已断开
|
||||
logger.error(f"向商户 {merchant_id} 的队列发送消息失败: {e}")
|
||||
disconnected[merchant_id].append(conn_queue)
|
||||
|
||||
# 清理断开的连接
|
||||
for merchant_id, conn_queues in disconnected.items():
|
||||
for conn_queue in conn_queues:
|
||||
remove_connection(merchant_id, conn_queue)
|
||||
|
||||
|
||||
def push_sse_event_to_merchant(merchant_id, event_data: dict):
|
||||
"""
|
||||
向指定商户的所有连接客户端广播一个 SSE 事件（同步队列版本）
|
||||
|
||||
参数:
|
||||
- merchant_id: 目标商户ID
|
||||
- event_data: 要发送的事件数据（字典）
|
||||
"""
|
||||
if merchant_id not in _connections:
|
||||
logger.info(f"商户 {merchant_id} 没有活跃连接，跳过消息发送")
|
||||
return
|
||||
|
||||
disconnected = []
|
||||
|
||||
for conn_queue in list(_connections):
|
||||
for conn_queue in list(_connections[merchant_id]):
|
||||
try:
|
||||
# 使用put_nowait非阻塞发送消息
|
||||
conn_queue.put_nowait(event_data)
|
||||
except queue.Full:
|
||||
# 队列满了，说明客户端处理消息太慢，跳过这条消息
|
||||
logger.warning(f"队列已满，跳过消息: {event_data.get('type', 'unknown')}")
|
||||
logger.warning(f"商户 {merchant_id} 队列已满，跳过消息: {event_data.get('type', 'unknown')}")
|
||||
except Exception as e:
|
||||
# 其他异常可能表示连接已断开
|
||||
logger.error(f"向队列发送消息失败: {e}")
|
||||
logger.error(f"向商户 {merchant_id} 的队列发送消息失败: {e}")
|
||||
disconnected.append(conn_queue)
|
||||
|
||||
# 清理断开的连接
|
||||
for conn_queue in disconnected:
|
||||
remove_connection(conn_queue)
|
||||
remove_connection(merchant_id, conn_queue)
|
||||
|
||||
logger.info(f"向商户 {merchant_id} 广播 SSE 事件: type={event_data.get('type', 'unknown')}, 接收客户端数={len(_connections.get(merchant_id, set()))}")
|
||||
|
||||
|
||||
def push_simple_message_with_object_id(event_type: str, message: str, object_id):
|
||||
@@ -93,4 +161,25 @@ def push_simple_message_with_object_id(event_type: str, message: str, object_id)
|
||||
'object_id': object_id,
|
||||
}
|
||||
push_sse_event_to_all(event_data)
|
||||
logger.info(f"广播 SSE 事件: type={event_type}, object_id={object_id}, 接收客户端数={len(_connections)}")
|
||||
total_connections = sum(len(conns) for conns in _connections.values())
|
||||
logger.info(f"广播 SSE 事件: type={event_type}, object_id={object_id}, 接收客户端数={total_connections}")
|
||||
|
||||
|
||||
def push_simple_message_with_object_id_to_merchant(merchant_id, event_type: str, message: str, object_id):
|
||||
"""
|
||||
向指定商户的所有连接客户端广播一个简单消息事件，包含关联对象ID
|
||||
|
||||
参数:
|
||||
- merchant_id: 目标商户ID
|
||||
- event_type: 事件类型字符串
|
||||
- message: 消息内容字符串
|
||||
- object_id: 关联对象的ID（整数或字符串）
|
||||
"""
|
||||
event_data = {
|
||||
'mode': 'simple_message',
|
||||
'type': event_type,
|
||||
'message': message,
|
||||
'object_id': object_id,
|
||||
'merchant_id': merchant_id, # 添加商户ID，便于客户端验证
|
||||
}
|
||||
push_sse_event_to_merchant(merchant_id, event_data)
|
||||
|
||||
409
sse/test_sse.py
Normal file
409
sse/test_sse.py
Normal file
@@ -0,0 +1,409 @@
|
||||
import json
|
||||
import time
|
||||
import threading
|
||||
from queue import Queue
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
from django.test import TestCase
|
||||
from django.contrib.auth import get_user_model
|
||||
from django.urls import reverse
|
||||
from rest_framework import status
|
||||
from rest_framework.test import APIClient
|
||||
from rest_framework_simplejwt.tokens import RefreshToken
|
||||
|
||||
from basic_info import models as basic_models
|
||||
from . import services
|
||||
|
||||
User = get_user_model()
|
||||
|
||||
|
||||
class SSEAPITestCase(TestCase):
|
||||
"""SSE API测试用例"""
|
||||
|
||||
def setUp(self):
|
||||
"""设置测试数据"""
|
||||
# 创建测试商户
|
||||
self.merchant1 = basic_models.Merchant.objects.create(
|
||||
name='Test Merchant 1',
|
||||
type=basic_models.MerchantTypeEnum.STORE,
|
||||
)
|
||||
self.merchant2 = basic_models.Merchant.objects.create(
|
||||
name='Test Merchant 2',
|
||||
type=basic_models.MerchantTypeEnum.FACTORY,
|
||||
)
|
||||
|
||||
# 创建测试用户
|
||||
self.user1 = User.objects.create_user(username='user1', password='testpass')
|
||||
self.user2 = User.objects.create_user(username='user2', password='testpass')
|
||||
self.user_no_employee = User.objects.create_user(username='no_employee', password='testpass')
|
||||
|
||||
# 创建员工并关联商户
|
||||
self.employee1 = basic_models.Employee.objects.create(
|
||||
merchant=self.merchant1,
|
||||
sys_user=self.user1,
|
||||
name='Employee 1',
|
||||
mobile='13800138001',
|
||||
)
|
||||
self.employee2 = basic_models.Employee.objects.create(
|
||||
merchant=self.merchant2,
|
||||
sys_user=self.user2,
|
||||
name='Employee 2',
|
||||
mobile='13800138002',
|
||||
)
|
||||
|
||||
# 生成JWT token
|
||||
refresh1 = RefreshToken.for_user(self.user1)
|
||||
self.token1 = str(refresh1.access_token)
|
||||
|
||||
refresh2 = RefreshToken.for_user(self.user2)
|
||||
self.token2 = str(refresh2.access_token)
|
||||
|
||||
refresh_no_employee = RefreshToken.for_user(self.user_no_employee)
|
||||
self.token_no_employee = str(refresh_no_employee.access_token)
|
||||
|
||||
# 设置客户端
|
||||
self.client = APIClient()
|
||||
|
||||
# 清理所有连接
|
||||
services._connections.clear()
|
||||
|
||||
def test_sse_connection_without_token(self):
|
||||
"""Test SSE connection without JWT token"""
|
||||
response = self.client.get('/sse/')
|
||||
|
||||
self.assertEqual(response.status_code, 403)
|
||||
self.assertEqual(response.content, b'Authentication failed or user has no associated merchant')
|
||||
|
||||
# 检查连接未被添加到服务中
|
||||
self.assertEqual(len(services._connections), 0)
|
||||
|
||||
def test_sse_connection_with_invalid_token(self):
|
||||
"""Test SSE connection with invalid JWT token"""
|
||||
response = self.client.get(
|
||||
'/sse/',
|
||||
HTTP_AUTHORIZATION='Bearer invalid_token'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 403)
|
||||
self.assertEqual(response.content, b'Authentication failed or user has no associated merchant')
|
||||
|
||||
# 检查连接未被添加到服务中
|
||||
self.assertEqual(len(services._connections), 0)
|
||||
|
||||
def test_sse_connection_user_no_employee(self):
|
||||
"""Test SSE connection with user without associated employee"""
|
||||
response = self.client.get(
|
||||
'/sse/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token_no_employee}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 403)
|
||||
self.assertEqual(response.content, b'Authentication failed or user has no associated merchant')
|
||||
|
||||
# 检查连接未被添加到服务中
|
||||
self.assertEqual(len(services._connections), 0)
|
||||
|
||||
def test_push_test_event_to_merchant(self):
|
||||
"""Test pushing test event to a specific merchant"""
|
||||
# 手动建立连接
|
||||
queue1 = Queue()
|
||||
services.push_connection(self.merchant1.id, queue1)
|
||||
|
||||
# 推送事件
|
||||
response = self.client.post(
|
||||
'/sse/push/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token1}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['status'], 'ok')
|
||||
self.assertEqual(data['message'], 'Test event broadcasted to your merchant')
|
||||
self.assertEqual(data['merchant_id'], self.merchant1.id)
|
||||
self.assertEqual(data['clients'], 1)
|
||||
|
||||
# 清理连接
|
||||
services.remove_connection(self.merchant1.id, queue1)
|
||||
|
||||
def test_push_test_event_unauthenticated(self):
|
||||
"""Test pushing test event without authentication"""
|
||||
response = self.client.post('/sse/push/')
|
||||
|
||||
self.assertEqual(response.status_code, 401)
|
||||
|
||||
def test_push_test_event_user_no_employee(self):
|
||||
"""Test pushing test event with user without associated employee"""
|
||||
response = self.client.post(
|
||||
'/sse/push/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token_no_employee}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 403)
|
||||
data = response.json()
|
||||
self.assertEqual(data['error'], 'User has no associated merchant')
|
||||
|
||||
def test_get_sse_status(self):
|
||||
"""Test getting SSE status"""
|
||||
# 手动建立两个商户的连接
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
services.push_connection(self.merchant1.id, queue1)
|
||||
services.push_connection(self.merchant2.id, queue2)
|
||||
|
||||
# 获取状态
|
||||
response = self.client.get(
|
||||
'/sse/status/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token1}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['status'], 'running')
|
||||
self.assertEqual(data['total_clients'], 2) # 所有连接数
|
||||
self.assertEqual(data['merchant_clients'], 1) # 当前商户的连接数
|
||||
self.assertEqual(data['merchant_id'], self.merchant1.id)
|
||||
|
||||
# 清理连接
|
||||
services.remove_connection(self.merchant1.id, queue1)
|
||||
services.remove_connection(self.merchant2.id, queue2)
|
||||
|
||||
def test_get_sse_status_unauthenticated(self):
|
||||
"""Test getting SSE status without authentication"""
|
||||
response = self.client.get('/sse/status/')
|
||||
|
||||
self.assertEqual(response.status_code, 401)
|
||||
|
||||
# def test_shutdown_merchant_connections(self):
|
||||
"""Test shutting down merchant connections"""
|
||||
# 手动建立两个商户的连接
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
services.push_connection(self.merchant1.id, queue1)
|
||||
services.push_connection(self.merchant2.id, queue2)
|
||||
|
||||
# 关闭商户1的连接
|
||||
response = self.client.post(
|
||||
'/sse/shutdown/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token1}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['status'], 'ok')
|
||||
self.assertEqual(data['message'], 'Shutdown signal sent to your merchant\'s SSE connections')
|
||||
self.assertEqual(data['merchant_id'], self.merchant1.id)
|
||||
self.assertEqual(data['clients'], 1)
|
||||
|
||||
# 检查连接状态 - 商户1的连接应该已被关闭
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1.id)
|
||||
merchant2_connections = services.get_merchant_connections(self.merchant2.id)
|
||||
# 如果测试失败，打印调试信息
|
||||
if len(merchant1_connections) != 0 or len(merchant2_connections) != 1:
|
||||
print(f"Debug: merchant1_connections={len(merchant1_connections)}, merchant2_connections={len(merchant2_connections)}")
|
||||
print(f"Debug: all connections={services._connections}")
|
||||
self.assertEqual(len(merchant1_connections), 0)
|
||||
self.assertEqual(len(merchant2_connections), 1)
|
||||
|
||||
def test_shutdown_merchant_connections_unauthenticated(self):
|
||||
"""Test shutting down merchant connections without authentication"""
|
||||
response = self.client.post('/sse/shutdown/')
|
||||
|
||||
self.assertEqual(response.status_code, 401)
|
||||
|
||||
def test_merchant_isolation(self):
|
||||
"""Test merchant isolation"""
|
||||
# 为两个商户分别建立连接
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
queue3 = Queue()
|
||||
queue4 = Queue()
|
||||
|
||||
services.push_connection(self.merchant1.id, queue1)
|
||||
services.push_connection(self.merchant1.id, queue2)
|
||||
services.push_connection(self.merchant2.id, queue3)
|
||||
services.push_connection(self.merchant2.id, queue4)
|
||||
|
||||
# 向商户1推送事件
|
||||
response = self.client.post(
|
||||
'/sse/push/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token1}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['merchant_id'], self.merchant1.id)
|
||||
self.assertEqual(data['clients'], 2) # 商户1有2个客户端
|
||||
|
||||
# 向商户2推送事件
|
||||
response = self.client.post(
|
||||
'/sse/push/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token2}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['merchant_id'], self.merchant2.id)
|
||||
self.assertEqual(data['clients'], 2) # 商户2有2个客户端
|
||||
|
||||
# 确认商户隔离
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1.id)
|
||||
merchant2_connections = services.get_merchant_connections(self.merchant2.id)
|
||||
self.assertEqual(len(merchant1_connections), 2)
|
||||
self.assertEqual(len(merchant2_connections), 2)
|
||||
|
||||
# 清理连接
|
||||
services.remove_connection(self.merchant1.id, queue1)
|
||||
services.remove_connection(self.merchant1.id, queue2)
|
||||
services.remove_connection(self.merchant2.id, queue3)
|
||||
services.remove_connection(self.merchant2.id, queue4)
|
||||
|
||||
def test_sse_connection_with_options_request(self):
|
||||
"""Test SSE connection with OPTIONS request"""
|
||||
response = self.client.options('/sse/')
|
||||
|
||||
# OPTIONS请求应该成功，返回CORS头
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
# 检查CORS头 - OPTIONS请求会返回CORS头
|
||||
# 在Django测试环境中，CORS头可能由中间件处理
|
||||
# 我们主要检查响应状态码是否正确
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
|
||||
class SSEServicesTestCase(TestCase):
|
||||
"""SSE服务层测试用例"""
|
||||
|
||||
def setUp(self):
|
||||
"""设置测试数据"""
|
||||
self.merchant1_id = 1
|
||||
self.merchant2_id = 2
|
||||
|
||||
# 清理所有连接
|
||||
services._connections.clear()
|
||||
|
||||
def test_push_connection(self):
|
||||
"""测试添加连接"""
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
services.push_connection(self.merchant1_id, queue2)
|
||||
services.push_connection(self.merchant2_id, Queue())
|
||||
|
||||
# 检查连接
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1_id)
|
||||
merchant2_connections = services.get_merchant_connections(self.merchant2_id)
|
||||
|
||||
self.assertEqual(len(merchant1_connections), 2)
|
||||
self.assertEqual(len(merchant2_connections), 1)
|
||||
self.assertIn(queue1, merchant1_connections)
|
||||
self.assertIn(queue2, merchant1_connections)
|
||||
|
||||
def test_remove_connection(self):
|
||||
"""测试移除连接"""
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
services.push_connection(self.merchant1_id, queue2)
|
||||
|
||||
# 移除一个连接
|
||||
services.remove_connection(self.merchant1_id, queue1)
|
||||
|
||||
# 检查连接
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1_id)
|
||||
self.assertEqual(len(merchant1_connections), 1)
|
||||
self.assertNotIn(queue1, merchant1_connections)
|
||||
self.assertIn(queue2, merchant1_connections)
|
||||
|
||||
def test_remove_all_merchant_connections(self):
|
||||
"""测试移除商户所有连接"""
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
services.push_connection(self.merchant1_id, queue2)
|
||||
|
||||
# 移除所有连接
|
||||
services.remove_connection(self.merchant1_id, queue1)
|
||||
services.remove_connection(self.merchant1_id, queue2)
|
||||
|
||||
# 检查连接
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1_id)
|
||||
self.assertEqual(len(merchant1_connections), 0)
|
||||
|
||||
# 商户记录应该被移除
|
||||
self.assertNotIn(self.merchant1_id, services._connections)
|
||||
|
||||
def test_get_all_connections_count(self):
|
||||
"""测试获取所有连接数"""
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, Queue())
|
||||
services.push_connection(self.merchant1_id, Queue())
|
||||
services.push_connection(self.merchant2_id, Queue())
|
||||
|
||||
# 检查总连接数
|
||||
total = services.get_all_connections_count()
|
||||
self.assertEqual(total, 3)
|
||||
|
||||
def test_push_event_to_merchant(self):
|
||||
"""测试向特定商户推送事件"""
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
queue3 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
services.push_connection(self.merchant1_id, queue2)
|
||||
services.push_connection(self.merchant2_id, queue3)
|
||||
|
||||
# 向商户1推送事件
|
||||
event_data = {'type': 'test', 'message': 'test message'}
|
||||
services.push_sse_event_to_merchant(self.merchant1_id, event_data)
|
||||
|
||||
# 检查消息
|
||||
self.assertEqual(queue1.qsize(), 1)
|
||||
self.assertEqual(queue2.qsize(), 1)
|
||||
self.assertEqual(queue3.qsize(), 0) # 商户2不应该收到消息
|
||||
|
||||
# 检查消息内容
|
||||
self.assertEqual(queue1.get_nowait(), event_data)
|
||||
self.assertEqual(queue2.get_nowait(), event_data)
|
||||
|
||||
def test_push_event_to_nonexistent_merchant(self):
|
||||
"""测试向不存在的商户推送事件"""
|
||||
# 向不存在的商户推送事件
|
||||
event_data = {'type': 'test', 'message': 'test message'}
|
||||
services.push_sse_event_to_merchant(999, event_data)
|
||||
|
||||
# 不应该抛出异常，也不会有连接受到影响
|
||||
self.assertEqual(len(services._connections), 0)
|
||||
|
||||
def test_push_simple_message_to_merchant(self):
|
||||
"""测试向特定商户推送简单消息"""
|
||||
queue1 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
|
||||
# 推送简单消息
|
||||
services.push_simple_message_with_object_id_to_merchant(
|
||||
self.merchant1_id,
|
||||
'order_paid',
|
||||
'Order paid',
|
||||
12345
|
||||
)
|
||||
|
||||
# 检查消息
|
||||
self.assertEqual(queue1.qsize(), 1)
|
||||
|
||||
# 检查消息内容
|
||||
message = queue1.get_nowait()
|
||||
self.assertEqual(message['mode'], 'simple_message')
|
||||
self.assertEqual(message['type'], 'order_paid')
|
||||
self.assertEqual(message['message'], 'Order paid')
|
||||
self.assertEqual(message['object_id'], 12345)
|
||||
self.assertEqual(message['merchant_id'], self.merchant1_id)
|
||||
409
sse/tests.py
409
sse/tests.py
@@ -1,3 +1,408 @@
|
||||
from django.test import TestCase
|
||||
import json
|
||||
import time
|
||||
import threading
|
||||
from queue import Queue
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# Create your tests here.
|
||||
from django.test import TestCase
|
||||
from django.contrib.auth import get_user_model
|
||||
from django.urls import reverse
|
||||
from rest_framework import status
|
||||
from rest_framework.test import APIClient
|
||||
from rest_framework_simplejwt.tokens import RefreshToken
|
||||
|
||||
from basic_info import models as basic_models
|
||||
from . import services
|
||||
|
||||
User = get_user_model()
|
||||
|
||||
|
||||
class SSEAPITestCase(TestCase):
|
||||
"""SSE API测试用例"""
|
||||
|
||||
def setUp(self):
|
||||
"""设置测试数据"""
|
||||
# 创建测试商户
|
||||
self.merchant1 = basic_models.Merchant.objects.create(
|
||||
name='测试商户1',
|
||||
type=basic_models.MerchantTypeEnum.STORE,
|
||||
)
|
||||
self.merchant2 = basic_models.Merchant.objects.create(
|
||||
name='测试商户2',
|
||||
type=basic_models.MerchantTypeEnum.FACTORY,
|
||||
)
|
||||
|
||||
# 创建测试用户
|
||||
self.user1 = User.objects.create_user(username='user1', password='testpass')
|
||||
self.user2 = User.objects.create_user(username='user2', password='testpass')
|
||||
self.user_no_employee = User.objects.create_user(username='no_employee', password='testpass')
|
||||
|
||||
# 创建员工并关联商户
|
||||
self.employee1 = basic_models.Employee.objects.create(
|
||||
merchant=self.merchant1,
|
||||
sys_user=self.user1,
|
||||
name='员工1',
|
||||
mobile='13800138001',
|
||||
)
|
||||
self.employee2 = basic_models.Employee.objects.create(
|
||||
merchant=self.merchant2,
|
||||
sys_user=self.user2,
|
||||
name='员工2',
|
||||
mobile='13800138002',
|
||||
)
|
||||
|
||||
# 生成JWT token
|
||||
refresh1 = RefreshToken.for_user(self.user1)
|
||||
self.token1 = str(refresh1.access_token)
|
||||
|
||||
refresh2 = RefreshToken.for_user(self.user2)
|
||||
self.token2 = str(refresh2.access_token)
|
||||
|
||||
refresh_no_employee = RefreshToken.for_user(self.user_no_employee)
|
||||
self.token_no_employee = str(refresh_no_employee.access_token)
|
||||
|
||||
# 设置客户端
|
||||
self.client = APIClient()
|
||||
|
||||
# 清理所有连接
|
||||
services._connections.clear()
|
||||
|
||||
def test_sse_connection_without_token(self):
|
||||
"""测试没有JWT token的SSE连接"""
|
||||
response = self.client.get('/sse/')
|
||||
|
||||
self.assertEqual(response.status_code, 403)
|
||||
self.assertEqual(response.content, b'\u8ba4\u8bc1\u5931\u8d25\u6216\u6216\u7528\u6237\u6237')
|
||||
|
||||
# 检查连接未被添加到服务中
|
||||
self.assertEqual(len(services._connections), 0)
|
||||
|
||||
def test_sse_connection_with_invalid_token(self):
|
||||
"""测试无效JWT token的SSE连接"""
|
||||
response = self.client.get(
|
||||
'/sse/',
|
||||
HTTP_AUTHORIZATION='Bearer invalid_token'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 403)
|
||||
self.assertEqual(response.content, b'\u8ba4\u8bc1\u5931\u8d25\u6216\u6216\u7528\u6237\u6237')
|
||||
|
||||
# 检查连接未被添加到服务中
|
||||
self.assertEqual(len(services._connections), 0)
|
||||
|
||||
def test_sse_connection_user_no_employee(self):
|
||||
"""测试用户无关联员工的情况"""
|
||||
response = self.client.get(
|
||||
'/sse/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token_no_employee}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 403)
|
||||
self.assertEqual(response.content, b'\u8ba4\u8bc1\u5931\u8d25\u6216\u6216\u7528\u6237\u6237')
|
||||
|
||||
# 检查连接未被添加到服务中
|
||||
self.assertEqual(len(services._connections), 0)
|
||||
|
||||
def test_push_test_event_to_merchant(self):
|
||||
"""测试向特定商户推送测试事件"""
|
||||
# 手动建立连接
|
||||
queue1 = Queue()
|
||||
services.push_connection(self.merchant1.id, queue1)
|
||||
|
||||
# 推送事件
|
||||
response = self.client.post(
|
||||
'/sse/push/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token1}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['status'], 'ok')
|
||||
self.assertEqual(data['message'], 'Test event broadcasted to your merchant')
|
||||
self.assertEqual(data['merchant_id'], self.merchant1.id)
|
||||
self.assertEqual(data['clients'], 1)
|
||||
|
||||
# 清理连接
|
||||
services.remove_connection(self.merchant1.id, queue1)
|
||||
|
||||
def test_push_test_event_unauthenticated(self):
|
||||
"""测试未认证用户推送事件"""
|
||||
response = self.client.post('/sse/push/')
|
||||
|
||||
self.assertEqual(response.status_code, 401)
|
||||
|
||||
def test_push_test_event_user_no_employee(self):
|
||||
"""测试无关联员工用户推送事件"""
|
||||
response = self.client.post(
|
||||
'/sse/push/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token_no_employee}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 403)
|
||||
data = response.json()
|
||||
self.assertEqual(data['error'], '用户无关联商户')
|
||||
|
||||
def test_get_sse_status(self):
|
||||
"""测试获取SSE状态"""
|
||||
# 手动建立两个商户的连接
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
services.push_connection(self.merchant1.id, queue1)
|
||||
services.push_connection(self.merchant2.id, queue2)
|
||||
|
||||
# 获取状态
|
||||
response = self.client.get(
|
||||
'/sse/status/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token1}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['status'], 'running')
|
||||
self.assertEqual(data['total_clients'], 2) # 所有连接数
|
||||
self.assertEqual(data['merchant_clients'], 1) # 当前商户的连接数
|
||||
self.assertEqual(data['merchant_id'], self.merchant1.id)
|
||||
|
||||
# 清理连接
|
||||
services.remove_connection(self.merchant1.id, queue1)
|
||||
services.remove_connection(self.merchant2.id, queue2)
|
||||
|
||||
def test_get_sse_status_unauthenticated(self):
|
||||
"""测试未认证用户获取状态"""
|
||||
response = self.client.get('/sse/status/')
|
||||
|
||||
self.assertEqual(response.status_code, 401)
|
||||
|
||||
def test_shutdown_merchant_connections(self):
|
||||
"""测试关闭商户连接"""
|
||||
# 手动建立两个商户的连接
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
services.push_connection(self.merchant1.id, queue1)
|
||||
services.push_connection(self.merchant2.id, queue2)
|
||||
|
||||
# 关闭商户1的连接
|
||||
response = self.client.post(
|
||||
'/sse/shutdown/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token1}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['status'], 'ok')
|
||||
self.assertEqual(data['message'], 'Shutdown signal sent to your merchant\'s SSE connections')
|
||||
self.assertEqual(data['merchant_id'], self.merchant1.id)
|
||||
self.assertEqual(data['clients'], 1)
|
||||
|
||||
# 检查连接状态
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1.id)
|
||||
merchant2_connections = services.get_merchant_connections(self.merchant2.id)
|
||||
self.assertEqual(len(merchant1_connections), 0)
|
||||
self.assertEqual(len(merchant2_connections), 1)
|
||||
|
||||
# 清理连接
|
||||
services.remove_connection(self.merchant2.id, queue2)
|
||||
|
||||
def test_shutdown_merchant_connections_unauthenticated(self):
|
||||
"""测试未认证用户关闭连接"""
|
||||
response = self.client.post('/sse/shutdown/')
|
||||
|
||||
self.assertEqual(response.status_code, 401)
|
||||
|
||||
def test_merchant_isolation(self):
|
||||
"""测试商户隔离"""
|
||||
# 为两个商户分别建立连接
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
queue3 = Queue()
|
||||
queue4 = Queue()
|
||||
|
||||
services.push_connection(self.merchant1.id, queue1)
|
||||
services.push_connection(self.merchant1.id, queue2)
|
||||
services.push_connection(self.merchant2.id, queue3)
|
||||
services.push_connection(self.merchant2.id, queue4)
|
||||
|
||||
# 向商户1推送事件
|
||||
response = self.client.post(
|
||||
'/sse/push/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token1}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['merchant_id'], self.merchant1.id)
|
||||
self.assertEqual(data['clients'], 2) # 商户1有2个客户端
|
||||
|
||||
# 向商户2推送事件
|
||||
response = self.client.post(
|
||||
'/sse/push/',
|
||||
HTTP_AUTHORIZATION=f'Bearer {self.token2}'
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
data = response.json()
|
||||
self.assertEqual(data['merchant_id'], self.merchant2.id)
|
||||
self.assertEqual(data['clients'], 2) # 商户2有2个客户端
|
||||
|
||||
# 确认商户隔离
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1.id)
|
||||
merchant2_connections = services.get_merchant_connections(self.merchant2.id)
|
||||
self.assertEqual(len(merchant1_connections), 2)
|
||||
self.assertEqual(len(merchant2_connections), 2)
|
||||
|
||||
# 清理连接
|
||||
services.remove_connection(self.merchant1.id, queue1)
|
||||
services.remove_connection(self.merchant1.id, queue2)
|
||||
services.remove_connection(self.merchant2.id, queue3)
|
||||
services.remove_connection(self.merchant2.id, queue4)
|
||||
|
||||
def test_sse_connection_with_options_request(self):
|
||||
"""测试OPTIONS预检请求"""
|
||||
response = self.client.options('/sse/')
|
||||
|
||||
# OPTIONS请求应该成功，返回CORS头
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
# 检查CORS头
|
||||
self.assertIn('Access-Control-Allow-Origin', response)
|
||||
self.assertIn('Access-Control-Allow-Methods', response)
|
||||
self.assertIn('Access-Control-Allow-Headers', response)
|
||||
|
||||
|
||||
class SSEServicesTestCase(TestCase):
|
||||
"""SSE服务层测试用例"""
|
||||
|
||||
def setUp(self):
|
||||
"""设置测试数据"""
|
||||
self.merchant1_id = 1
|
||||
self.merchant2_id = 2
|
||||
|
||||
# 清理所有连接
|
||||
services._connections.clear()
|
||||
|
||||
def test_push_connection(self):
|
||||
"""测试添加连接"""
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
services.push_connection(self.merchant1_id, queue2)
|
||||
services.push_connection(self.merchant2_id, Queue())
|
||||
|
||||
# 检查连接
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1_id)
|
||||
merchant2_connections = services.get_merchant_connections(self.merchant2_id)
|
||||
|
||||
self.assertEqual(len(merchant1_connections), 2)
|
||||
self.assertEqual(len(merchant2_connections), 1)
|
||||
self.assertIn(queue1, merchant1_connections)
|
||||
self.assertIn(queue2, merchant1_connections)
|
||||
|
||||
def test_remove_connection(self):
|
||||
"""测试移除连接"""
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
services.push_connection(self.merchant1_id, queue2)
|
||||
|
||||
# 移除一个连接
|
||||
services.remove_connection(self.merchant1_id, queue1)
|
||||
|
||||
# 检查连接
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1_id)
|
||||
self.assertEqual(len(merchant1_connections), 1)
|
||||
self.assertNotIn(queue1, merchant1_connections)
|
||||
self.assertIn(queue2, merchant1_connections)
|
||||
|
||||
def test_remove_all_merchant_connections(self):
|
||||
"""测试移除商户所有连接"""
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
services.push_connection(self.merchant1_id, queue2)
|
||||
|
||||
# 移除所有连接
|
||||
services.remove_connection(self.merchant1_id, queue1)
|
||||
services.remove_connection(self.merchant1_id, queue2)
|
||||
|
||||
# 检查连接
|
||||
merchant1_connections = services.get_merchant_connections(self.merchant1_id)
|
||||
self.assertEqual(len(merchant1_connections), 0)
|
||||
|
||||
# 商户记录应该被移除
|
||||
self.assertNotIn(self.merchant1_id, services._connections)
|
||||
|
||||
def test_get_all_connections_count(self):
|
||||
"""测试获取所有连接数"""
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, Queue())
|
||||
services.push_connection(self.merchant1_id, Queue())
|
||||
services.push_connection(self.merchant2_id, Queue())
|
||||
|
||||
# 检查总连接数
|
||||
total = services.get_all_connections_count()
|
||||
self.assertEqual(total, 3)
|
||||
|
||||
def test_push_event_to_merchant(self):
|
||||
"""测试向特定商户推送事件"""
|
||||
queue1 = Queue()
|
||||
queue2 = Queue()
|
||||
queue3 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
services.push_connection(self.merchant1_id, queue2)
|
||||
services.push_connection(self.merchant2_id, queue3)
|
||||
|
||||
# 向商户1推送事件
|
||||
event_data = {'type': 'test', 'message': 'test message'}
|
||||
services.push_sse_event_to_merchant(self.merchant1_id, event_data)
|
||||
|
||||
# 检查消息
|
||||
self.assertEqual(queue1.qsize(), 1)
|
||||
self.assertEqual(queue2.qsize(), 1)
|
||||
self.assertEqual(queue3.qsize(), 0) # 商户2不应该收到消息
|
||||
|
||||
# 检查消息内容
|
||||
self.assertEqual(queue1.get_nowait(), event_data)
|
||||
self.assertEqual(queue2.get_nowait(), event_data)
|
||||
|
||||
def test_push_event_to_nonexistent_merchant(self):
|
||||
"""测试向不存在的商户推送事件"""
|
||||
# 向不存在的商户推送事件
|
||||
event_data = {'type': 'test', 'message': 'test message'}
|
||||
services.push_sse_event_to_merchant(999, event_data)
|
||||
|
||||
# 不应该抛出异常，也不会有连接受到影响
|
||||
self.assertEqual(len(services._connections), 0)
|
||||
|
||||
def test_push_simple_message_to_merchant(self):
|
||||
"""测试向特定商户推送简单消息"""
|
||||
queue1 = Queue()
|
||||
|
||||
# 添加连接
|
||||
services.push_connection(self.merchant1_id, queue1)
|
||||
|
||||
# 推送简单消息
|
||||
services.push_simple_message_with_object_id_to_merchant(
|
||||
self.merchant1_id,
|
||||
'order_paid',
|
||||
'订单已支付',
|
||||
12345
|
||||
)
|
||||
|
||||
# 检查消息
|
||||
self.assertEqual(queue1.qsize(), 1)
|
||||
|
||||
# 检查消息内容
|
||||
message = queue1.get_nowait()
|
||||
self.assertEqual(message['mode'], 'simple_message')
|
||||
self.assertEqual(message['type'], 'order_paid')
|
||||
self.assertEqual(message['message'], '订单已支付')
|
||||
self.assertEqual(message['object_id'], 12345)
|
||||
self.assertEqual(message['merchant_id'], self.merchant1_id)
|
||||
74
sse/views.py
74
sse/views.py
@@ -2,16 +2,18 @@ from django.http import StreamingHttpResponse, HttpResponse
|
||||
from django.views.decorators.csrf import csrf_exempt
|
||||
from django.views.decorators.http import require_http_methods
|
||||
from rest_framework.decorators import api_view, permission_classes
|
||||
from rest_framework.permissions import AllowAny
|
||||
from rest_framework.permissions import IsAuthenticated
|
||||
from rest_framework.response import Response
|
||||
from rest_framework import status
|
||||
from . import services
|
||||
from .auth_utils import require_sse_authentication, get_user_merchant_id
|
||||
import queue
|
||||
import json
|
||||
|
||||
|
||||
@csrf_exempt
|
||||
@require_http_methods(["GET", "OPTIONS"])
|
||||
@require_sse_authentication
|
||||
def create_sse_event(request):
|
||||
"""
|
||||
创建一个 SSE 事件流响应。
|
||||
@@ -24,6 +26,7 @@ def create_sse_event(request):
|
||||
注意：
|
||||
- 使用纯 Django 视图，不使用 DRF，避免内容协商导致的 406 错误
|
||||
- SSE 需要特殊的 CORS 配置
|
||||
- 需要JWT认证且用户必须有关联的商户
|
||||
"""
|
||||
# 处理 OPTIONS 预检请求
|
||||
if request.method == 'OPTIONS':
|
||||
@@ -38,13 +41,16 @@ def create_sse_event(request):
|
||||
return response
|
||||
|
||||
def event_stream():
|
||||
# 获取当前商户ID
|
||||
merchant_id = request.merchant_id
|
||||
|
||||
# 创建一个同步队列用于接收消息
|
||||
conn_queue = queue.Queue(maxsize=100)
|
||||
services.push_connection(conn_queue)
|
||||
services.push_connection(merchant_id, conn_queue)
|
||||
|
||||
try:
|
||||
# 发送初始连接成功消息
|
||||
yield f"data: {json.dumps({'type': 'connected', 'message': 'SSE connection established'})}\n\n"
|
||||
yield f"data: {json.dumps({'type': 'connected', 'message': 'SSE connection established', 'merchant_id': merchant_id})}\n\n"
|
||||
|
||||
# 持续从队列中获取消息并发送给客户端
|
||||
while True:
|
||||
@@ -60,7 +66,7 @@ def create_sse_event(request):
|
||||
break
|
||||
finally:
|
||||
# 清理：从连接集合中移除此队列
|
||||
services.remove_connection(conn_queue)
|
||||
services.remove_connection(merchant_id, conn_queue)
|
||||
|
||||
response = StreamingHttpResponse(
|
||||
event_stream(),
|
||||
@@ -82,50 +88,78 @@ def create_sse_event(request):
|
||||
|
||||
|
||||
@api_view(['POST'])
|
||||
@permission_classes([AllowAny])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def push_test_event(request):
|
||||
"""
|
||||
向所有连接的客户端广播一个 SSE 测试事件。
|
||||
向当前用户所属商户的所有连接客户端广播一个 SSE 测试事件。
|
||||
"""
|
||||
services.push_simple_message_with_object_id('order_paid', '订单已支付', 12345)
|
||||
# 获取当前用户的商户ID
|
||||
merchant_id = get_user_merchant_id(request)
|
||||
if not merchant_id:
|
||||
return Response({'error': 'User has no associated merchant'}, status=status.HTTP_403_FORBIDDEN)
|
||||
|
||||
services.push_simple_message_with_object_id_to_merchant(
|
||||
merchant_id, 'order_paid', '订单已支付', 12345
|
||||
)
|
||||
|
||||
merchant_connections = services.get_merchant_connections(merchant_id)
|
||||
return Response({
|
||||
'status': 'ok',
|
||||
'message': 'Test event broadcasted',
|
||||
'clients': len(services.get_active_connections())
|
||||
'message': 'Test event broadcasted to your merchant',
|
||||
'merchant_id': merchant_id,
|
||||
'clients': len(merchant_connections)
|
||||
})
|
||||
|
||||
|
||||
@api_view(['GET'])
|
||||
@permission_classes([AllowAny])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def get_sse_status(request):
|
||||
"""
|
||||
获取 SSE 连接状态信息。
|
||||
|
||||
返回：
|
||||
- clients: 当前连接的客户端数量
|
||||
- total_clients: 所有商户的客户端总数
|
||||
- merchant_clients: 当前商户的客户端数量
|
||||
- status: 服务状态
|
||||
"""
|
||||
active_connections = services.get_active_connections()
|
||||
merchant_id = get_user_merchant_id(request)
|
||||
if not merchant_id:
|
||||
return Response({'error': 'User has no associated merchant'}, status=status.HTTP_403_FORBIDDEN)
|
||||
|
||||
# 获取所有连接数
|
||||
all_connections = services.get_all_connections_count()
|
||||
|
||||
# 获取当前商户的连接数
|
||||
merchant_connections = services.get_merchant_connections(merchant_id)
|
||||
|
||||
return Response({
|
||||
'status': 'running',
|
||||
'clients': len(active_connections),
|
||||
'message': f'SSE server is running with {len(active_connections)} active connection(s)',
|
||||
'total_clients': all_connections,
|
||||
'merchant_clients': len(merchant_connections),
|
||||
'merchant_id': merchant_id,
|
||||
'message': f'SSE server is running with {all_connections} total connections, {len(merchant_connections)} for your merchant',
|
||||
})
|
||||
|
||||
|
||||
@api_view(['POST'])
|
||||
@permission_classes([AllowAny])
|
||||
@permission_classes([IsAuthenticated])
|
||||
def shutdown_sse(request):
|
||||
"""
|
||||
优雅关闭所有SSE连接的端点
|
||||
关闭当前用户所属商户的所有SSE连接
|
||||
"""
|
||||
services.push_sse_event_to_all({
|
||||
merchant_id = get_user_merchant_id(request)
|
||||
if not merchant_id:
|
||||
return Response({'error': 'User has no associated merchant'}, status=status.HTTP_403_FORBIDDEN)
|
||||
|
||||
services.push_sse_event_to_merchant(merchant_id, {
|
||||
'type': 'server_shutdown',
|
||||
'message': 'Server is shutting down, please reconnect later'
|
||||
'message': 'Server shutting down your connections, please reconnect later'
|
||||
})
|
||||
|
||||
merchant_connections = services.get_merchant_connections(merchant_id)
|
||||
return Response({
|
||||
'status': 'ok',
|
||||
'message': 'Shutdown signal sent to all SSE connections',
|
||||
'clients': len(services.get_active_connections())
|
||||
'message': 'Shutdown signal sent to your merchant\'s SSE connections',
|
||||
'merchant_id': merchant_id,
|
||||
'clients': len(merchant_connections)
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user