585 lines
21 KiB
Python
585 lines
21 KiB
Python
"""
|
|
任务管理器测试用例
|
|
|
|
覆盖场景：
|
|
1. 创建检测任务（按需、定期、批量）
|
|
2. 查询任务列表
|
|
3. 执行检测任务
|
|
4. 任务状态管理
|
|
5. 错误处理和重试机制
|
|
"""
|
|
import pytest
|
|
from datetime import datetime, date, timedelta
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy import select, and_
|
|
|
|
from app.services.risk_detection.task_manager import TaskManager
|
|
from app.services.risk_detection.task_manager.scheduler import DetectionScheduler
|
|
from app.models.risk_detection import (
|
|
DetectionTask,
|
|
DetectionRule,
|
|
TaskType,
|
|
TaskStatus,
|
|
RiskLevel,
|
|
)
|
|
from app.models.risk_detection import TaskExecution
|
|
|
|
|
|
class TestTaskManager:
|
|
"""任务管理器测试类"""
|
|
|
|
@pytest.fixture
|
|
def mock_db_session(self):
|
|
"""创建Mock数据库会话"""
|
|
session = AsyncMock(spec=AsyncSession)
|
|
session.execute = AsyncMock()
|
|
session.flush = AsyncMock()
|
|
session.refresh = AsyncMock()
|
|
session.add = MagicMock()
|
|
session.commit = AsyncMock()
|
|
return session
|
|
|
|
@pytest.fixture
|
|
def task_manager(self, mock_db_session):
|
|
"""创建任务管理器实例"""
|
|
return TaskManager(mock_db_session)
|
|
|
|
@pytest.fixture
|
|
def mock_detection_rules(self):
|
|
"""Mock检测规则数据"""
|
|
return [
|
|
DetectionRule(
|
|
rule_id="rule_revenue_001",
|
|
rule_name="收入完整性检测",
|
|
algorithm_code="REVENUE_INTEGRITY_CHECK",
|
|
description="检测平台充值与申报收入的匹配度",
|
|
is_enabled=True,
|
|
parameters={"threshold": 0.1},
|
|
created_at=datetime.now(),
|
|
),
|
|
DetectionRule(
|
|
rule_id="rule_private_001",
|
|
rule_name="私户收款检测",
|
|
algorithm_code="PRIVATE_ACCOUNT_DETECTION",
|
|
description="识别使用私人账户收款的风险",
|
|
is_enabled=True,
|
|
parameters={"threshold_amount": 10000},
|
|
created_at=datetime.now(),
|
|
),
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_task_on_demand(self, task_manager, mock_db_session, mock_detection_rules):
|
|
"""测试创建按需检测任务"""
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalar_one_or_none.return_value = None
|
|
mock_db_session.execute.return_value.scalars.return_value.all.return_value = mock_detection_rules
|
|
|
|
# 创建任务
|
|
task = await task_manager.create_task(
|
|
task_name="测试按需检测任务",
|
|
task_type=TaskType.ON_DEMAND,
|
|
entity_ids=["ZB_TEST_001"],
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
rule_ids=["rule_revenue_001"],
|
|
parameters={"description": "测试任务"}
|
|
)
|
|
|
|
# 验证结果
|
|
assert task.task_id is not None
|
|
assert task.task_name == "测试按需检测任务"
|
|
assert task.task_type == TaskType.ON_DEMAND
|
|
assert task.status == TaskStatus.PENDING
|
|
assert task.entity_type == "streamer"
|
|
assert task.period == "2024-01"
|
|
assert task.total_entities == 1
|
|
assert task.parameters["description"] == "测试任务"
|
|
|
|
# 验证执行记录被创建
|
|
mock_db_session.add.assert_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_task_periodic(self, task_manager, mock_db_session, mock_detection_rules):
|
|
"""测试创建定期检测任务"""
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalar_one_or_none.return_value = None
|
|
mock_db_session.execute.return_value.scalars.return_value.all.return_value = mock_detection_rules
|
|
|
|
# 创建任务
|
|
task = await task_manager.create_task(
|
|
task_name="测试定期检测任务",
|
|
task_type=TaskType.PERIODIC,
|
|
entity_ids=["ZB_TEST_001", "ZB_TEST_002"],
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
parameters={"schedule_type": "periodic"}
|
|
)
|
|
|
|
# 验证结果
|
|
assert task.task_id is not None
|
|
assert task.task_name == "测试定期检测任务"
|
|
assert task.task_type == TaskType.PERIODIC
|
|
assert task.total_entities == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_task_batch(self, task_manager, mock_db_session, mock_detection_rules):
|
|
"""测试创建批量检测任务"""
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalar_one_or_none.return_value = None
|
|
mock_db_session.execute.return_value.scalars.return_value.all.return_value = mock_detection_rules
|
|
|
|
# 创建任务
|
|
task = await task_manager.create_task(
|
|
task_name="测试批量检测任务",
|
|
task_type=TaskType.BATCH,
|
|
entity_ids=["ZB_TEST_001", "ZB_TEST_002", "ZB_TEST_003"],
|
|
entity_type="mcn",
|
|
period="2024-01",
|
|
rule_ids=["rule_private_001"],
|
|
parameters={"batch_size": 3}
|
|
)
|
|
|
|
# 验证结果
|
|
assert task.task_id is not None
|
|
assert task.task_name == "测试批量检测任务"
|
|
assert task.task_type == TaskType.BATCH
|
|
assert task.total_entities == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_task(self, task_manager, mock_db_session):
|
|
"""测试获取任务"""
|
|
# 创建mock任务
|
|
mock_task = DetectionTask(
|
|
task_id="task_test_001",
|
|
task_name="测试任务",
|
|
task_type=TaskType.ON_DEMAND,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=1,
|
|
processed_entities=0,
|
|
parameters={},
|
|
created_at=datetime.now(),
|
|
)
|
|
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalar_one_or_none.return_value = mock_task
|
|
|
|
# 执行测试
|
|
task = await task_manager.get_task("task_test_001")
|
|
|
|
# 验证结果
|
|
assert task is not None
|
|
assert task.task_id == "task_test_001"
|
|
assert task.task_name == "测试任务"
|
|
assert task.status == TaskStatus.PENDING
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_task_not_found(self, task_manager, mock_db_session):
|
|
"""测试获取不存在的任务"""
|
|
# 设置mock返回值为None
|
|
mock_db_session.execute.return_value.scalar_one_or_none.return_value = None
|
|
|
|
# 执行测试
|
|
task = await task_manager.get_task("task_not_exist")
|
|
|
|
# 验证结果
|
|
assert task is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_tasks(self, task_manager, mock_db_session):
|
|
"""测试查询任务列表"""
|
|
# 创建mock任务列表
|
|
mock_tasks = [
|
|
DetectionTask(
|
|
task_id="task_001",
|
|
task_name="任务1",
|
|
task_type=TaskType.ON_DEMAND,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=1,
|
|
processed_entities=0,
|
|
parameters={},
|
|
created_at=datetime.now(),
|
|
),
|
|
DetectionTask(
|
|
task_id="task_002",
|
|
task_name="任务2",
|
|
task_type=TaskType.PERIODIC,
|
|
status=TaskStatus.COMPLETED,
|
|
entity_type="mcn",
|
|
period="2024-01",
|
|
total_entities=5,
|
|
processed_entities=5,
|
|
parameters={},
|
|
created_at=datetime.now(),
|
|
),
|
|
]
|
|
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalars.return_value.all.return_value = mock_tasks
|
|
|
|
# 执行测试 - 无过滤条件
|
|
tasks = await task_manager.list_tasks()
|
|
assert len(tasks) == 2
|
|
|
|
# 执行测试 - 按状态过滤
|
|
tasks = await task_manager.list_tasks(status=TaskStatus.PENDING)
|
|
assert len(tasks) == 2 # 实际项目中这里应该过滤
|
|
|
|
# 执行测试 - 按类型过滤
|
|
tasks = await task_manager.list_tasks(task_type=TaskType.ON_DEMAND)
|
|
assert len(tasks) == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_task_success(self, task_manager, mock_db_session):
|
|
"""测试成功执行任务"""
|
|
# 创建mock任务
|
|
mock_task = DetectionTask(
|
|
task_id="task_test_001",
|
|
task_name="测试任务",
|
|
task_type=TaskType.ON_DEMAND,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=1,
|
|
processed_entities=0,
|
|
parameters={},
|
|
created_at=datetime.now(),
|
|
)
|
|
|
|
# 创建mock任务执行记录
|
|
mock_executions = [
|
|
TaskExecution(
|
|
execution_id="exec_001",
|
|
task_id="task_test_001",
|
|
entity_id="ZB_TEST_001",
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
rule_ids="rule_revenue_001",
|
|
parameters={},
|
|
status="pending",
|
|
created_at=datetime.now(),
|
|
)
|
|
]
|
|
|
|
# 设置mock返回值
|
|
mock_db_session.execute.side_effect = [
|
|
AsyncMock(scalar_one_or_none=lambda: mock_task), # 获取任务
|
|
AsyncMock(scalars=lambda: AsyncMock(all=lambda: mock_executions)), # 获取执行记录
|
|
AsyncMock(scalars=lambda: AsyncMock(all=lambda: [])), # 获取规则
|
|
]
|
|
|
|
# Mock RuleEngine
|
|
with patch('app.services.risk_detection.task_manager.task_manager.RuleEngine') as MockRuleEngine:
|
|
mock_rule_engine = AsyncMock()
|
|
mock_rule_engine.execute_detection = AsyncMock(return_value={
|
|
"results": [
|
|
{
|
|
"entity_id": "ZB_TEST_001",
|
|
"risk_level": RiskLevel.MEDIUM,
|
|
"risk_score": 65.0,
|
|
"description": "发现收入不匹配问题",
|
|
"suggestion": "请检查收入申报是否准确",
|
|
}
|
|
],
|
|
"summary": {
|
|
"total_detections": 1,
|
|
"risk_distribution": {"MEDIUM": 1},
|
|
"avg_score": 65.0,
|
|
}
|
|
})
|
|
MockRuleEngine.return_value = mock_rule_engine
|
|
|
|
# Mock算法注册
|
|
with patch('app.services.risk_detection.task_manager.task_manager.RevenueIntegrityAlgorithm'):
|
|
# 执行任务
|
|
result = await task_manager.execute_task("task_test_001")
|
|
|
|
# 验证结果
|
|
assert result["task_id"] == "task_test_001"
|
|
assert result["status"] == "completed"
|
|
assert result["summary"]["total_entities"] == 1
|
|
assert result["summary"]["total_detections"] == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_task_not_found(self, task_manager, mock_db_session):
|
|
"""测试执行不存在的任务"""
|
|
# 设置mock返回值为None
|
|
mock_db_session.execute.return_value.scalar_one_or_none.return_value = None
|
|
|
|
# 验证抛出异常
|
|
with pytest.raises(ValueError, match="任务不存在"):
|
|
await task_manager.execute_task("task_not_exist")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_task_with_failure(self, task_manager, mock_db_session):
|
|
"""测试任务执行失败"""
|
|
# 创建mock任务
|
|
mock_task = DetectionTask(
|
|
task_id="task_test_001",
|
|
task_name="测试任务",
|
|
task_type=TaskType.ON_DEMAND,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=1,
|
|
processed_entities=0,
|
|
parameters={},
|
|
created_at=datetime.now(),
|
|
)
|
|
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalar_one_or_none.return_value = mock_task
|
|
|
|
# Mock RuleEngine抛出异常
|
|
with patch('app.services.risk_detection.task_manager.task_manager.RuleEngine') as MockRuleEngine:
|
|
mock_rule_engine = AsyncMock()
|
|
mock_rule_engine.execute_detection = AsyncMock(side_effect=Exception("执行失败"))
|
|
MockRuleEngine.return_value = mock_rule_engine
|
|
|
|
# 验证抛出异常
|
|
with pytest.raises(Exception, match="执行失败"):
|
|
await task_manager.execute_task("task_test_001")
|
|
|
|
|
|
class TestDetectionScheduler:
|
|
"""任务调度器测试类"""
|
|
|
|
@pytest.fixture
|
|
def mock_db_session(self):
|
|
"""创建Mock数据库会话"""
|
|
session = AsyncMock(spec=AsyncSession)
|
|
session.execute = AsyncMock()
|
|
return session
|
|
|
|
@pytest.fixture
|
|
def scheduler(self, mock_db_session):
|
|
"""创建调度器实例"""
|
|
return DetectionScheduler(mock_db_session)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_periodic_detection(self, scheduler, mock_db_session):
|
|
"""测试调度定期检测"""
|
|
# Mock TaskManager.create_task
|
|
mock_task = DetectionTask(
|
|
task_id="task_periodic_001",
|
|
task_name="定期检测-2024-01",
|
|
task_type=TaskType.PERIODIC,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=2,
|
|
processed_entities=0,
|
|
parameters={"schedule_type": "periodic"},
|
|
created_at=datetime.now(),
|
|
)
|
|
|
|
with patch.object(scheduler.task_manager, 'create_task', new_callable=AsyncMock) as mock_create:
|
|
mock_create.return_value = mock_task
|
|
|
|
# 执行测试
|
|
task = await scheduler.schedule_periodic_detection(
|
|
entity_ids=["ZB_TEST_001", "ZB_TEST_002"],
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
auto_execute=False
|
|
)
|
|
|
|
# 验证结果
|
|
assert task.task_id == "task_periodic_001"
|
|
assert task.task_type == TaskType.PERIODIC
|
|
mock_create.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_on_demand_detection(self, scheduler, mock_db_session):
|
|
"""测试调度按需检测"""
|
|
# Mock TaskManager.create_task
|
|
mock_task = DetectionTask(
|
|
task_id="task_ondemand_001",
|
|
task_name="按需检测-ZB_TEST_001-2024-01",
|
|
task_type=TaskType.ON_DEMAND,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=1,
|
|
processed_entities=0,
|
|
parameters={"schedule_type": "on_demand"},
|
|
created_at=datetime.now(),
|
|
)
|
|
|
|
with patch.object(scheduler.task_manager, 'create_task', new_callable=AsyncMock) as mock_create:
|
|
mock_create.return_value = mock_task
|
|
|
|
# 执行测试
|
|
task = await scheduler.schedule_on_demand_detection(
|
|
entity_id="ZB_TEST_001",
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
auto_execute=False
|
|
)
|
|
|
|
# 验证结果
|
|
assert task.task_id == "task_ondemand_001"
|
|
assert task.task_type == TaskType.ON_DEMAND
|
|
mock_create.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_batch_detection(self, scheduler, mock_db_session):
|
|
"""测试调度批量检测"""
|
|
# Mock TaskManager.create_task
|
|
mock_task = DetectionTask(
|
|
task_id="task_batch_001",
|
|
task_name="批量检测-2024-01",
|
|
task_type=TaskType.BATCH,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=10,
|
|
processed_entities=0,
|
|
parameters={"schedule_type": "batch", "batch_size": 10},
|
|
created_at=datetime.now(),
|
|
)
|
|
|
|
with patch.object(scheduler.task_manager, 'create_task', new_callable=AsyncMock) as mock_create:
|
|
mock_create.return_value = mock_task
|
|
|
|
# 执行测试
|
|
task = await scheduler.schedule_batch_detection(
|
|
entity_ids=[f"ZB_TEST_{i:03d}" for i in range(1, 11)],
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
auto_execute=False
|
|
)
|
|
|
|
# 验证结果
|
|
assert task.task_id == "task_batch_001"
|
|
assert task.task_type == TaskType.BATCH
|
|
assert task.total_entities == 10
|
|
mock_create.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_pending_tasks(self, scheduler, mock_db_session):
|
|
"""测试检查待执行任务"""
|
|
# 创建mock待执行任务
|
|
mock_tasks = [
|
|
DetectionTask(
|
|
task_id="task_001",
|
|
task_name="待执行任务1",
|
|
task_type=TaskType.PERIODIC,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=1,
|
|
processed_entities=0,
|
|
parameters={},
|
|
created_at=datetime.now(),
|
|
),
|
|
DetectionTask(
|
|
task_id="task_002",
|
|
task_name="待执行任务2",
|
|
task_type=TaskType.BATCH,
|
|
status=TaskStatus.PENDING,
|
|
entity_type="mcn",
|
|
period="2024-01",
|
|
total_entities=5,
|
|
processed_entities=0,
|
|
parameters={},
|
|
created_at=datetime.now(),
|
|
),
|
|
]
|
|
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalars.return_value.all.return_value = mock_tasks
|
|
|
|
# 执行测试
|
|
tasks = await scheduler.check_pending_tasks()
|
|
|
|
# 验证结果
|
|
assert len(tasks) == 2
|
|
assert all(task.status == TaskStatus.PENDING for task in tasks)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retry_failed_tasks(self, scheduler, mock_db_session):
|
|
"""测试重试失败任务"""
|
|
# 创建mock失败任务
|
|
mock_tasks = [
|
|
DetectionTask(
|
|
task_id="task_failed_001",
|
|
task_name="失败任务1",
|
|
task_type=TaskType.ON_DEMAND,
|
|
status=TaskStatus.FAILED,
|
|
entity_type="streamer",
|
|
period="2024-01",
|
|
total_entities=1,
|
|
processed_entities=0,
|
|
retry_count=1,
|
|
parameters={},
|
|
created_at=datetime.now(),
|
|
)
|
|
]
|
|
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalars.return_value.all.return_value = mock_tasks
|
|
|
|
# Mock任务执行
|
|
with patch.object(scheduler, '_execute_task_async', new_callable=AsyncMock):
|
|
# 执行测试
|
|
retried_tasks = await scheduler.retry_failed_tasks(max_retries=3)
|
|
|
|
# 验证结果
|
|
assert len(retried_tasks) == 1
|
|
assert retried_tasks[0].status == TaskStatus.PENDING
|
|
assert retried_tasks[0].retry_count == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_old_tasks(self, scheduler, mock_db_session):
|
|
"""测试清理旧任务"""
|
|
# 创建mock旧任务
|
|
old_date = datetime.now() - timedelta(days=35)
|
|
mock_tasks = [
|
|
DetectionTask(
|
|
task_id="task_old_001",
|
|
task_name="旧任务1",
|
|
task_type=TaskType.ON_DEMAND,
|
|
status=TaskStatus.COMPLETED,
|
|
entity_type="streamer",
|
|
period="2023-01",
|
|
total_entities=1,
|
|
processed_entities=1,
|
|
parameters={},
|
|
created_at=old_date,
|
|
),
|
|
DetectionTask(
|
|
task_id="task_old_002",
|
|
task_name="旧任务2",
|
|
task_type=TaskType.ON_DEMAND,
|
|
status=TaskStatus.FAILED,
|
|
entity_type="mcn",
|
|
period="2023-01",
|
|
total_entities=1,
|
|
processed_entities=0,
|
|
parameters={},
|
|
created_at=old_date,
|
|
),
|
|
]
|
|
|
|
# 设置mock返回值
|
|
mock_db_session.execute.return_value.scalars.return_value.all.return_value = mock_tasks
|
|
mock_db_session.delete = MagicMock()
|
|
|
|
# 执行测试
|
|
count = await scheduler.cleanup_old_tasks(days=30)
|
|
|
|
# 验证结果
|
|
assert count == 2
|
|
assert mock_db_session.delete.call_count == 2
|
|
mock_db_session.commit.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|