forked from erp-dev/erp
409 lines
15 KiB
Python
409 lines
15 KiB
Python
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) |