538 lines
16 KiB
Python
538 lines
16 KiB
Python
"""
|
||
执行计划器单元测试
|
||
"""
|
||
import pytest
|
||
from unittest.mock import Mock, MagicMock
|
||
from datetime import datetime
|
||
|
||
from app.models.risk_detection import DetectionRule
|
||
from app.services.risk_detection.engine.execution_plan import (
|
||
ExecutionMode,
|
||
ExecutionNode,
|
||
ExecutionStage,
|
||
ExecutionPlan,
|
||
ExecutionPlanner,
|
||
ExecutionPlanBuilder,
|
||
)
|
||
|
||
|
||
class TestExecutionMode:
|
||
"""执行模式枚举测试"""
|
||
|
||
def test_execution_mode_values(self):
|
||
"""测试执行模式值"""
|
||
assert ExecutionMode.SEQUENTIAL == "sequential"
|
||
assert ExecutionMode.PARALLEL == "parallel"
|
||
assert ExecutionMode.HYBRID == "hybrid"
|
||
|
||
|
||
class TestExecutionNode:
|
||
"""执行节点测试"""
|
||
|
||
def test_init(self):
|
||
"""测试节点初始化"""
|
||
rule = Mock(spec=DetectionRule)
|
||
rule.rule_id = "rule_1"
|
||
|
||
node = ExecutionNode(
|
||
node_id="node_1",
|
||
rule_id="rule_1",
|
||
rule=rule,
|
||
level=0,
|
||
dependencies=["node_0"],
|
||
dependents=["node_2"]
|
||
)
|
||
|
||
assert node.node_id == "node_1"
|
||
assert node.rule_id == "rule_1"
|
||
assert node.rule == rule
|
||
assert node.level == 0
|
||
assert node.dependencies == ["node_0"]
|
||
assert node.dependents == ["node_2"]
|
||
|
||
def test_to_dict(self):
|
||
"""测试转换为字典"""
|
||
rule = Mock(spec=DetectionRule)
|
||
rule.rule_id = "rule_1"
|
||
rule.rule_name = "规则1"
|
||
rule.algorithm_code = "ALGO_1"
|
||
|
||
node = ExecutionNode(
|
||
node_id="node_1",
|
||
rule_id="rule_1",
|
||
rule=rule,
|
||
level=1,
|
||
dependencies=["node_0"],
|
||
dependents=["node_2"]
|
||
)
|
||
|
||
result = node.to_dict()
|
||
|
||
assert result["node_id"] == "node_1"
|
||
assert result["rule_id"] == "rule_1"
|
||
assert result["algorithm_code"] == "ALGO_1"
|
||
assert result["algorithm_name"] == "规则1"
|
||
assert result["level"] == 1
|
||
assert result["dependencies"] == ["node_0"]
|
||
assert result["dependents"] == ["node_2"]
|
||
|
||
|
||
class TestExecutionStage:
|
||
"""执行阶段测试"""
|
||
|
||
def setup_method(self):
|
||
"""每个测试方法前执行"""
|
||
self.rule = Mock(spec=DetectionRule)
|
||
self.rule.rule_id = "rule_1"
|
||
|
||
self.node = ExecutionNode(
|
||
node_id="node_1",
|
||
rule_id="rule_1",
|
||
rule=self.rule,
|
||
level=0
|
||
)
|
||
|
||
def test_init(self):
|
||
"""测试阶段初始化"""
|
||
stage = ExecutionStage(
|
||
stage_id="stage_1",
|
||
stage_name="阶段1",
|
||
nodes=[self.node],
|
||
execution_mode=ExecutionMode.PARALLEL,
|
||
level=0,
|
||
depends_on=["stage_0"]
|
||
)
|
||
|
||
assert stage.stage_id == "stage_1"
|
||
assert stage.stage_name == "阶段1"
|
||
assert stage.nodes == [self.node]
|
||
assert stage.execution_mode == ExecutionMode.PARALLEL
|
||
assert stage.level == 0
|
||
assert stage.depends_on == ["stage_0"]
|
||
|
||
def test_rules_property(self):
|
||
"""测试rules属性"""
|
||
stage = ExecutionStage(
|
||
stage_id="stage_1",
|
||
stage_name="阶段1",
|
||
nodes=[self.node],
|
||
execution_mode=ExecutionMode.PARALLEL
|
||
)
|
||
|
||
assert stage.rules == [self.rule]
|
||
|
||
def test_rule_count_property(self):
|
||
"""测试rule_count属性"""
|
||
stage = ExecutionStage(
|
||
stage_id="stage_1",
|
||
stage_name="阶段1",
|
||
nodes=[self.node],
|
||
execution_mode=ExecutionMode.PARALLEL
|
||
)
|
||
|
||
assert stage.rule_count == 1
|
||
|
||
def test_to_dict(self):
|
||
"""测试转换为字典"""
|
||
stage = ExecutionStage(
|
||
stage_id="stage_1",
|
||
stage_name="阶段1",
|
||
nodes=[self.node],
|
||
execution_mode=ExecutionMode.PARALLEL,
|
||
level=0,
|
||
depends_on=["stage_0"]
|
||
)
|
||
|
||
result = stage.to_dict()
|
||
|
||
assert result["stage_id"] == "stage_1"
|
||
assert result["stage_name"] == "阶段1"
|
||
assert result["level"] == 0
|
||
assert result["execution_mode"] == "parallel"
|
||
assert result["rule_count"] == 1
|
||
assert result["depends_on"] == ["stage_0"]
|
||
assert len(result["nodes"]) == 1
|
||
|
||
|
||
class TestExecutionPlan:
|
||
"""执行计划测试"""
|
||
|
||
def setup_method(self):
|
||
"""每个测试方法前执行"""
|
||
self.rule = Mock(spec=DetectionRule)
|
||
self.rule.rule_id = "rule_1"
|
||
|
||
self.node = ExecutionNode(
|
||
node_id="node_1",
|
||
rule_id="rule_1",
|
||
rule=self.rule,
|
||
level=0
|
||
)
|
||
|
||
self.stage = ExecutionStage(
|
||
stage_id="stage_1",
|
||
stage_name="阶段1",
|
||
nodes=[self.node],
|
||
execution_mode=ExecutionMode.PARALLEL
|
||
)
|
||
|
||
def test_init(self):
|
||
"""测试计划初始化"""
|
||
plan = ExecutionPlan(
|
||
plan_id="plan_1",
|
||
task_id="task_1",
|
||
entity_id="entity_1",
|
||
entity_type="type_1",
|
||
period="2024-01",
|
||
stages=[self.stage],
|
||
execution_mode=ExecutionMode.HYBRID,
|
||
total_rules=1,
|
||
max_level=0
|
||
)
|
||
|
||
assert plan.plan_id == "plan_1"
|
||
assert plan.task_id == "task_1"
|
||
assert plan.entity_id == "entity_1"
|
||
assert plan.entity_type == "type_1"
|
||
assert plan.period == "2024-01"
|
||
assert plan.stages == [self.stage]
|
||
assert plan.execution_mode == ExecutionMode.HYBRID
|
||
assert plan.total_rules == 1
|
||
assert plan.max_level == 0
|
||
|
||
def test_get_stage(self):
|
||
"""测试获取阶段"""
|
||
plan = ExecutionPlan(
|
||
plan_id="plan_1",
|
||
task_id="task_1",
|
||
entity_id="entity_1",
|
||
entity_type="type_1",
|
||
period="2024-01",
|
||
stages=[self.stage]
|
||
)
|
||
|
||
stage = plan.get_stage("stage_1")
|
||
assert stage == self.stage
|
||
|
||
stage = plan.get_stage("nonexistent")
|
||
assert stage is None
|
||
|
||
def test_get_stages_by_level(self):
|
||
"""测试按层级获取阶段"""
|
||
plan = ExecutionPlan(
|
||
plan_id="plan_1",
|
||
task_id="task_1",
|
||
entity_id="entity_1",
|
||
entity_type="type_1",
|
||
period="2024-01",
|
||
stages=[self.stage]
|
||
)
|
||
|
||
stages = plan.get_stages_by_level(0)
|
||
assert stages == [self.stage]
|
||
|
||
def test_get_next_stage_no_completed(self):
|
||
"""测试获取下一个阶段（无已完成）"""
|
||
plan = ExecutionPlan(
|
||
plan_id="plan_1",
|
||
task_id="task_1",
|
||
entity_id="entity_1",
|
||
entity_type="type_1",
|
||
period="2024-01",
|
||
stages=[self.stage]
|
||
)
|
||
|
||
next_stage = plan.get_next_stage([])
|
||
assert next_stage == self.stage
|
||
|
||
def test_get_next_stage_with_completed(self):
|
||
"""测试获取下一个阶段（有已完成）"""
|
||
stage2 = ExecutionStage(
|
||
stage_id="stage_2",
|
||
stage_name="阶段2",
|
||
nodes=[],
|
||
execution_mode=ExecutionMode.PARALLEL,
|
||
depends_on=["stage_1"]
|
||
)
|
||
|
||
plan = ExecutionPlan(
|
||
plan_id="plan_1",
|
||
task_id="task_1",
|
||
entity_id="entity_1",
|
||
entity_type="type_1",
|
||
period="2024-01",
|
||
stages=[self.stage, stage2]
|
||
)
|
||
|
||
# stage_1未完成，stage_2不能执行
|
||
next_stage = plan.get_next_stage([])
|
||
assert next_stage == self.stage
|
||
|
||
# stage_1已完成，stage_2可以执行
|
||
next_stage = plan.get_next_stage(["stage_1"])
|
||
assert next_stage == stage2
|
||
|
||
def test_get_execution_summary(self):
|
||
"""测试获取执行摘要"""
|
||
plan = ExecutionPlan(
|
||
plan_id="plan_1",
|
||
task_id="task_1",
|
||
entity_id="entity_1",
|
||
entity_type="type_1",
|
||
period="2024-01",
|
||
stages=[self.stage],
|
||
execution_mode=ExecutionMode.HYBRID,
|
||
total_rules=1,
|
||
max_level=0
|
||
)
|
||
|
||
summary = plan.get_execution_summary()
|
||
|
||
assert summary["plan_id"] == "plan_1"
|
||
assert summary["task_id"] == "task_1"
|
||
assert summary["entity_id"] == "entity_1"
|
||
assert summary["entity_type"] == "type_1"
|
||
assert summary["period"] == "2024-01"
|
||
assert summary["execution_mode"] == "hybrid"
|
||
assert summary["total_stages"] == 1
|
||
assert summary["total_rules"] == 1
|
||
assert summary["max_level"] == 0
|
||
assert summary["stages_per_level"] == {0: 1}
|
||
assert summary["max_parallelism"] == 1
|
||
|
||
def test_to_dict(self):
|
||
"""测试转换为字典"""
|
||
plan = ExecutionPlan(
|
||
plan_id="plan_1",
|
||
task_id="task_1",
|
||
entity_id="entity_1",
|
||
entity_type="type_1",
|
||
period="2024-01",
|
||
stages=[self.stage],
|
||
execution_mode=ExecutionMode.HYBRID,
|
||
total_rules=1,
|
||
max_level=0
|
||
)
|
||
|
||
result = plan.to_dict()
|
||
|
||
assert result["plan_id"] == "plan_1"
|
||
assert result["task_id"] == "task_1"
|
||
assert result["entity_id"] == "entity_1"
|
||
assert result["entity_type"] == "type_1"
|
||
assert result["period"] == "2024-01"
|
||
assert result["execution_mode"] == "hybrid"
|
||
assert result["total_stages"] == 1
|
||
assert result["total_rules"] == 1
|
||
assert result["max_level"] == 0
|
||
assert len(result["stages"]) == 1
|
||
assert "summary" in result
|
||
assert "created_at" in result
|
||
|
||
|
||
class TestExecutionPlanner:
|
||
"""执行计划器测试"""
|
||
|
||
def setup_method(self):
|
||
"""每个测试方法前执行"""
|
||
self.planner = ExecutionPlanner()
|
||
|
||
# 创建模拟规则
|
||
self.rule1 = Mock(spec=DetectionRule)
|
||
self.rule1.rule_id = "rule_1"
|
||
self.rule1.algorithm_code = "ALGO_1"
|
||
|
||
self.rule2 = Mock(spec=DetectionRule)
|
||
self.rule2.rule_id = "rule_2"
|
||
self.rule2.algorithm_code = "ALGO_2"
|
||
|
||
def test_init(self):
|
||
"""测试初始化"""
|
||
assert self.planner is not None
|
||
assert self.planner.dependency_resolver is not None
|
||
|
||
def test_create_sequential_plan(self):
|
||
"""测试创建串行执行计划"""
|
||
rules = [self.rule1, self.rule2]
|
||
|
||
plan = self.planner._create_sequential_plan(
|
||
"task_1", "entity_1", "type_1", "2024-01", rules
|
||
)
|
||
|
||
assert plan.task_id == "task_1"
|
||
assert plan.entity_id == "entity_1"
|
||
assert plan.entity_type == "type_1"
|
||
assert plan.period == "2024-01"
|
||
assert plan.execution_mode == ExecutionMode.SEQUENTIAL
|
||
assert plan.total_rules == 2
|
||
assert plan.max_level == 0
|
||
assert len(plan.stages) == 1
|
||
|
||
stage = plan.stages[0]
|
||
assert stage.execution_mode == ExecutionMode.SEQUENTIAL
|
||
assert stage.rule_count == 2
|
||
assert len(stage.nodes) == 2
|
||
|
||
def test_create_parallel_plan(self):
|
||
"""测试创建并行执行计划"""
|
||
rules = [self.rule1, self.rule2]
|
||
|
||
plan = self.planner._create_parallel_plan(
|
||
"task_1", "entity_1", "type_1", "2024-01", rules
|
||
)
|
||
|
||
assert plan.task_id == "task_1"
|
||
assert plan.execution_mode == ExecutionMode.PARALLEL
|
||
assert plan.total_rules == 2
|
||
assert len(plan.stages) == 1
|
||
|
||
stage = plan.stages[0]
|
||
assert stage.execution_mode == ExecutionMode.PARALLEL
|
||
assert len(stage.nodes) == 2
|
||
|
||
def test_create_hybrid_plan(self):
|
||
"""测试创建混合执行计划"""
|
||
rules = [self.rule1, self.rule2]
|
||
|
||
# 模拟无依赖的情况
|
||
plan = self.planner._create_hybrid_plan(
|
||
"task_1", "entity_1", "type_1", "2024-01", rules
|
||
)
|
||
|
||
assert plan.task_id == "task_1"
|
||
assert plan.execution_mode == ExecutionMode.HYBRID
|
||
assert plan.total_rules == 2
|
||
|
||
def test_create_plan_sequential_mode(self):
|
||
"""测试创建执行计划（串行模式）"""
|
||
rules = [self.rule1, self.rule2]
|
||
|
||
plan = self.planner.create_plan(
|
||
"task_1", "entity_1", "type_1", "2024-01", rules,
|
||
ExecutionMode.SEQUENTIAL
|
||
)
|
||
|
||
assert plan.execution_mode == ExecutionMode.SEQUENTIAL
|
||
|
||
def test_create_plan_parallel_mode(self):
|
||
"""测试创建执行计划（并行模式）"""
|
||
rules = [self.rule1, self.rule2]
|
||
|
||
plan = self.planner.create_plan(
|
||
"task_1", "entity_1", "type_1", "2024-01", rules,
|
||
ExecutionMode.PARALLEL
|
||
)
|
||
|
||
assert plan.execution_mode == ExecutionMode.PARALLEL
|
||
|
||
def test_create_plan_hybrid_mode(self):
|
||
"""测试创建执行计划（混合模式）"""
|
||
rules = [self.rule1, self.rule2]
|
||
|
||
plan = self.planner.create_plan(
|
||
"task_1", "entity_1", "type_1", "2024-01", rules,
|
||
ExecutionMode.HYBRID
|
||
)
|
||
|
||
assert plan.execution_mode == ExecutionMode.HYBRID
|
||
|
||
def test_optimize_plan(self):
|
||
"""测试优化执行计划"""
|
||
plan = ExecutionPlan(
|
||
plan_id="plan_1",
|
||
task_id="task_1",
|
||
entity_id="entity_1",
|
||
entity_type="type_1",
|
||
period="2024-01",
|
||
stages=[]
|
||
)
|
||
|
||
optimized = self.planner.optimize_plan(plan)
|
||
|
||
# 目前优化方法返回原计划
|
||
assert optimized == plan
|
||
|
||
|
||
class TestExecutionPlanBuilder:
|
||
"""执行计划构建器测试"""
|
||
|
||
def setup_method(self):
|
||
"""每个测试方法前执行"""
|
||
self.builder = ExecutionPlanBuilder(
|
||
"task_1", "entity_1", "type_1", "2024-01"
|
||
)
|
||
|
||
self.rule = Mock(spec=DetectionRule)
|
||
self.rule.rule_id = "rule_1"
|
||
|
||
self.node = ExecutionNode(
|
||
node_id="node_1",
|
||
rule_id="rule_1",
|
||
rule=self.rule
|
||
)
|
||
|
||
def test_init(self):
|
||
"""测试初始化"""
|
||
assert self.builder.task_id == "task_1"
|
||
assert self.builder.entity_id == "entity_1"
|
||
assert self.builder.entity_type == "type_1"
|
||
assert self.builder.period == "2024-01"
|
||
assert self.builder.stages == []
|
||
|
||
def test_add_stage(self):
|
||
"""测试添加阶段"""
|
||
result = self.builder.add_stage(
|
||
"阶段1",
|
||
[self.node],
|
||
ExecutionMode.PARALLEL,
|
||
[]
|
||
)
|
||
|
||
# 返回自身，支持链式调用
|
||
assert result == self.builder
|
||
|
||
# 阶段已添加
|
||
assert len(self.builder.stages) == 1
|
||
|
||
stage = self.builder.stages[0]
|
||
assert stage.stage_id == "stage_0"
|
||
assert stage.stage_name == "阶段1"
|
||
assert stage.nodes == [self.node]
|
||
assert stage.execution_mode == ExecutionMode.PARALLEL
|
||
assert stage.level == 0
|
||
assert stage.depends_on == []
|
||
|
||
def test_add_stage_with_dependencies(self):
|
||
"""测试添加带依赖的阶段"""
|
||
self.builder.add_stage("阶段1", [self.node])
|
||
|
||
self.builder.add_stage(
|
||
"阶段2",
|
||
[self.node],
|
||
ExecutionMode.PARALLEL,
|
||
["stage_0"]
|
||
)
|
||
|
||
assert len(self.builder.stages) == 2
|
||
|
||
stage2 = self.builder.stages[1]
|
||
assert stage2.stage_id == "stage_1"
|
||
assert stage2.level == 1
|
||
assert stage2.depends_on == ["stage_0"]
|
||
|
||
def test_build(self):
|
||
"""测试构建执行计划"""
|
||
self.builder.add_stage("阶段1", [self.node])
|
||
|
||
plan = self.builder.build()
|
||
|
||
assert plan.plan_id == "plan_task_1"
|
||
assert plan.task_id == "task_1"
|
||
assert plan.entity_id == "entity_1"
|
||
assert plan.entity_type == "type_1"
|
||
assert plan.period == "2024-01"
|
||
assert plan.execution_mode == ExecutionMode.HYBRID
|
||
assert plan.total_rules == 1
|
||
assert plan.max_level == 0
|
||
assert len(plan.stages) == 1
|