starmen1265
7ffaafd615
- Add tests for WorkflowState methods (mark_step_started, mark_step_completed, mark_step_failed) - Add tests for StepResult schema - Add tests for StepDefinition schema with dependencies - Achieve 100% test coverage on kb_pipeline/schemas.py - Total 37 tests passing Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
475 lines
16 KiB
Python
475 lines
16 KiB
Python
"""Schema 验证测试
|
|
|
|
测试 kb_pipeline.schemas 模块中的所有核心 Schema。
|
|
"""
|
|
|
|
import pytest
|
|
from datetime import datetime
|
|
from pydantic import ValidationError
|
|
|
|
|
|
class TestStepInput:
|
|
"""测试 StepInput Schema"""
|
|
|
|
def test_step_input_requires_content_id(self):
|
|
"""StepInput 必须包含 content_id 字段"""
|
|
from kb_pipeline.schemas import StepInput
|
|
|
|
with pytest.raises(ValidationError) as exc_info:
|
|
StepInput(platform="bilibili")
|
|
|
|
assert "content_id" in str(exc_info.value)
|
|
|
|
def test_step_input_requires_platform(self):
|
|
"""StepInput 必须包含 platform 字段"""
|
|
from kb_pipeline.schemas import StepInput
|
|
|
|
with pytest.raises(ValidationError) as exc_info:
|
|
StepInput(content_id="BV1test123")
|
|
|
|
assert "platform" in str(exc_info.value)
|
|
|
|
def test_step_input_validates_platform_literal(self):
|
|
"""StepInput platform 必须是 bilibili 或 zhihu"""
|
|
from kb_pipeline.schemas import StepInput
|
|
|
|
# 有效值
|
|
step_bilibili = StepInput(content_id="BV1test", platform="bilibili")
|
|
assert step_bilibili.platform == "bilibili"
|
|
|
|
step_zhihu = StepInput(content_id="123456", platform="zhihu")
|
|
assert step_zhihu.platform == "zhihu"
|
|
|
|
# 无效值
|
|
with pytest.raises(ValidationError):
|
|
StepInput(content_id="BV1test", platform="youtube")
|
|
|
|
def test_step_input_content_id_non_empty(self):
|
|
"""StepInput content_id 不能为空字符串"""
|
|
from kb_pipeline.schemas import StepInput
|
|
|
|
with pytest.raises(ValidationError):
|
|
StepInput(content_id="", platform="bilibili")
|
|
|
|
|
|
class TestStepOutput:
|
|
"""测试 StepOutput Schema"""
|
|
|
|
def test_step_output_requires_success(self):
|
|
"""StepOutput 必须包含 success 字段"""
|
|
from kb_pipeline.schemas import StepOutput
|
|
|
|
with pytest.raises(ValidationError) as exc_info:
|
|
StepOutput()
|
|
|
|
assert "success" in str(exc_info.value)
|
|
|
|
def test_step_output_default_completed_at(self):
|
|
"""StepOutput completed_at 默认为当前时间"""
|
|
from kb_pipeline.schemas import StepOutput
|
|
|
|
before = datetime.now()
|
|
output = StepOutput(success=True)
|
|
after = datetime.now()
|
|
|
|
assert before <= output.completed_at <= after
|
|
|
|
def test_step_output_optional_error_message(self):
|
|
"""StepOutput error_message 是可选的"""
|
|
from kb_pipeline.schemas import StepOutput
|
|
|
|
output = StepOutput(success=True)
|
|
assert output.error_message is None
|
|
|
|
output_with_error = StepOutput(success=False, error_message="测试错误")
|
|
assert output_with_error.error_message == "测试错误"
|
|
|
|
|
|
class TestStepStatus:
|
|
"""测试 StepStatus 枚举"""
|
|
|
|
def test_step_status_has_pending(self):
|
|
"""StepStatus 包含 PENDING 状态"""
|
|
from kb_pipeline.schemas import StepStatus
|
|
|
|
assert StepStatus.PENDING.value == "pending"
|
|
|
|
def test_step_status_has_running(self):
|
|
"""StepStatus 包含 RUNNING 状态"""
|
|
from kb_pipeline.schemas import StepStatus
|
|
|
|
assert StepStatus.RUNNING.value == "running"
|
|
|
|
def test_step_status_has_completed(self):
|
|
"""StepStatus 包含 COMPLETED 状态"""
|
|
from kb_pipeline.schemas import StepStatus
|
|
|
|
assert StepStatus.COMPLETED.value == "completed"
|
|
|
|
def test_step_status_has_failed(self):
|
|
"""StepStatus 包含 FAILED 状态"""
|
|
from kb_pipeline.schemas import StepStatus
|
|
|
|
assert StepStatus.FAILED.value == "failed"
|
|
|
|
def test_step_status_has_skipped(self):
|
|
"""StepStatus 包含 SKIPPED 状态"""
|
|
from kb_pipeline.schemas import StepStatus
|
|
|
|
assert StepStatus.SKIPPED.value == "skipped"
|
|
|
|
|
|
class TestStepRecord:
|
|
"""测试 StepRecord Schema"""
|
|
|
|
def test_step_record_requires_step_name(self):
|
|
"""StepRecord 必须包含 step_name"""
|
|
from kb_pipeline.schemas import StepRecord, StepStatus
|
|
|
|
with pytest.raises(ValidationError):
|
|
StepRecord(status=StepStatus.COMPLETED)
|
|
|
|
def test_step_record_requires_status(self):
|
|
"""StepRecord 必须包含 status"""
|
|
from kb_pipeline.schemas import StepRecord
|
|
|
|
with pytest.raises(ValidationError):
|
|
StepRecord(step_name="fetch_metadata")
|
|
|
|
def test_step_record_optional_fields(self):
|
|
"""StepRecord 可选字段正确工作"""
|
|
from kb_pipeline.schemas import StepRecord, StepStatus
|
|
|
|
record = StepRecord(step_name="fetch_metadata", status=StepStatus.COMPLETED)
|
|
assert record.started_at is None
|
|
assert record.completed_at is None
|
|
assert record.error_message is None
|
|
assert record.output_summary is None
|
|
|
|
|
|
class TestWorkflowState:
|
|
"""测试 WorkflowState Schema"""
|
|
|
|
def test_workflow_state_requires_content_id(self):
|
|
"""WorkflowState 必须包含 content_id"""
|
|
from kb_pipeline.schemas import WorkflowState
|
|
|
|
with pytest.raises(ValidationError):
|
|
WorkflowState(platform="bilibili", source_url="https://test.com")
|
|
|
|
def test_workflow_state_requires_platform(self):
|
|
"""WorkflowState 必须包含 platform"""
|
|
from kb_pipeline.schemas import WorkflowState
|
|
|
|
with pytest.raises(ValidationError):
|
|
WorkflowState(content_id="BV1test", source_url="https://test.com")
|
|
|
|
def test_workflow_state_requires_source_url(self):
|
|
"""WorkflowState 必须包含 source_url"""
|
|
from kb_pipeline.schemas import WorkflowState
|
|
|
|
with pytest.raises(ValidationError):
|
|
WorkflowState(content_id="BV1test", platform="bilibili")
|
|
|
|
def test_workflow_state_has_steps_dict(self):
|
|
"""WorkflowState 包含 steps 字典"""
|
|
from kb_pipeline.schemas import WorkflowState
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
assert hasattr(state, "steps")
|
|
assert isinstance(state.steps, dict)
|
|
|
|
def test_workflow_state_default_schema_version(self):
|
|
"""WorkflowState schema_version 默认为 1.0"""
|
|
from kb_pipeline.schemas import WorkflowState
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
assert state.schema_version == "1.0"
|
|
|
|
def test_workflow_state_is_step_completed_method(self):
|
|
"""WorkflowState 有 is_step_completed 方法"""
|
|
from kb_pipeline.schemas import WorkflowState, StepStatus, StepRecord
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
|
|
# 未完成的步骤
|
|
assert state.is_step_completed("fetch_metadata") is False
|
|
|
|
# 添加完成的步骤
|
|
state.steps["fetch_metadata"] = StepRecord(
|
|
step_name="fetch_metadata",
|
|
status=StepStatus.COMPLETED
|
|
)
|
|
assert state.is_step_completed("fetch_metadata") is True
|
|
|
|
|
|
class TestWorkflowStateMethods:
|
|
"""测试 WorkflowState 方法"""
|
|
|
|
def test_mark_step_started(self):
|
|
"""mark_step_started 正确设置步骤状态"""
|
|
from kb_pipeline.schemas import WorkflowState, StepStatus
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
|
|
state.mark_step_started("fetch_metadata")
|
|
|
|
assert "fetch_metadata" in state.steps
|
|
assert state.steps["fetch_metadata"].status == StepStatus.RUNNING
|
|
assert state.steps["fetch_metadata"].started_at is not None
|
|
|
|
def test_mark_step_completed_existing_step(self):
|
|
"""mark_step_completed 更新已存在的步骤"""
|
|
from kb_pipeline.schemas import WorkflowState, StepStatus
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
|
|
state.mark_step_started("fetch_metadata")
|
|
state.mark_step_completed("fetch_metadata", "成功获取元数据")
|
|
|
|
assert state.steps["fetch_metadata"].status == StepStatus.COMPLETED
|
|
assert state.steps["fetch_metadata"].completed_at is not None
|
|
assert state.steps["fetch_metadata"].output_summary == "成功获取元数据"
|
|
|
|
def test_mark_step_completed_new_step(self):
|
|
"""mark_step_completed 创建新步骤记录"""
|
|
from kb_pipeline.schemas import WorkflowState, StepStatus
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
|
|
state.mark_step_completed("download_video", "视频下载完成")
|
|
|
|
assert "download_video" in state.steps
|
|
assert state.steps["download_video"].status == StepStatus.COMPLETED
|
|
assert state.steps["download_video"].output_summary == "视频下载完成"
|
|
|
|
def test_mark_step_failed_existing_step(self):
|
|
"""mark_step_failed 更新已存在的步骤为失败"""
|
|
from kb_pipeline.schemas import WorkflowState, StepStatus
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
|
|
state.mark_step_started("fetch_metadata")
|
|
state.mark_step_failed("fetch_metadata", "网络错误")
|
|
|
|
assert state.steps["fetch_metadata"].status == StepStatus.FAILED
|
|
assert state.steps["fetch_metadata"].error_message == "网络错误"
|
|
assert state.steps["fetch_metadata"].completed_at is not None
|
|
|
|
def test_mark_step_failed_new_step(self):
|
|
"""mark_step_failed 创建新的失败步骤记录"""
|
|
from kb_pipeline.schemas import WorkflowState, StepStatus
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
|
|
state.mark_step_failed("download_video", "文件不存在")
|
|
|
|
assert "download_video" in state.steps
|
|
assert state.steps["download_video"].status == StepStatus.FAILED
|
|
assert state.steps["download_video"].error_message == "文件不存在"
|
|
|
|
|
|
class TestStepResult:
|
|
"""测试 StepResult Schema"""
|
|
|
|
def test_step_result_creation(self):
|
|
"""StepResult 正确创建"""
|
|
from kb_pipeline.schemas import StepResult, StepInput, StepStatus
|
|
|
|
input_data = StepInput(content_id="BV1test", platform="bilibili")
|
|
result = StepResult(step_name="fetch_metadata", input_data=input_data)
|
|
|
|
assert result.step_name == "fetch_metadata"
|
|
assert result.status == StepStatus.PENDING
|
|
assert result.output_data is None
|
|
|
|
def test_step_result_with_output(self):
|
|
"""StepResult 包含输出数据"""
|
|
from kb_pipeline.schemas import StepResult, StepInput, StepOutput, StepStatus
|
|
|
|
input_data = StepInput(content_id="BV1test", platform="bilibili")
|
|
output_data = StepOutput(success=True)
|
|
result = StepResult(
|
|
step_name="fetch_metadata",
|
|
input_data=input_data,
|
|
output_data=output_data,
|
|
status=StepStatus.COMPLETED
|
|
)
|
|
|
|
assert result.output_data is not None
|
|
assert result.output_data.success is True
|
|
|
|
|
|
class TestStepDefinition:
|
|
"""测试 StepDefinition Schema"""
|
|
|
|
def test_step_definition_creation(self):
|
|
"""StepDefinition 正确创建"""
|
|
from kb_pipeline.schemas import StepDefinition, StepInput, StepOutput
|
|
|
|
definition = StepDefinition(
|
|
step_name="fetch_metadata",
|
|
description="获取视频元数据",
|
|
input_schema=StepInput,
|
|
output_schema=StepOutput
|
|
)
|
|
|
|
assert definition.step_name == "fetch_metadata"
|
|
assert definition.description == "获取视频元数据"
|
|
assert definition.dependencies == []
|
|
|
|
def test_step_definition_with_dependencies(self):
|
|
"""StepDefinition 包含依赖"""
|
|
from kb_pipeline.schemas import StepDefinition, StepInput, StepOutput
|
|
|
|
definition = StepDefinition(
|
|
step_name="transcribe_audio",
|
|
description="转录音频",
|
|
input_schema=StepInput,
|
|
output_schema=StepOutput,
|
|
dependencies=["download_video", "extract_audio"]
|
|
)
|
|
|
|
assert len(definition.dependencies) == 2
|
|
assert "download_video" in definition.dependencies
|
|
|
|
|
|
class TestSchemaValidation:
|
|
"""测试 Pydantic 验证"""
|
|
|
|
def test_validation_error_on_invalid_input(self):
|
|
"""无效输入抛出 ValidationError"""
|
|
from kb_pipeline.schemas import StepInput
|
|
|
|
with pytest.raises(ValidationError):
|
|
StepInput(content_id=123, platform="bilibili") # type: ignore
|
|
|
|
def test_model_validate_from_dict(self):
|
|
"""model_validate 从字典正确创建实例"""
|
|
from kb_pipeline.schemas import StepInput
|
|
|
|
data = {"content_id": "BV1test", "platform": "bilibili"}
|
|
step_input = StepInput.model_validate(data)
|
|
assert step_input.content_id == "BV1test"
|
|
assert step_input.platform == "bilibili"
|
|
|
|
|
|
class TestJSONSerialization:
|
|
"""测试 JSON 序列化/反序列化"""
|
|
|
|
def test_step_input_model_dump_json(self):
|
|
"""StepInput 支持 model_dump_json"""
|
|
from kb_pipeline.schemas import StepInput
|
|
import json
|
|
|
|
step_input = StepInput(content_id="BV1test", platform="bilibili")
|
|
json_str = step_input.model_dump_json()
|
|
data = json.loads(json_str)
|
|
|
|
assert data["content_id"] == "BV1test"
|
|
assert data["platform"] == "bilibili"
|
|
|
|
def test_step_output_model_dump_json(self):
|
|
"""StepOutput 支持 model_dump_json"""
|
|
from kb_pipeline.schemas import StepOutput
|
|
import json
|
|
|
|
output = StepOutput(success=True, error_message=None)
|
|
json_str = output.model_dump_json()
|
|
data = json.loads(json_str)
|
|
|
|
assert data["success"] is True
|
|
assert data["error_message"] is None
|
|
|
|
def test_workflow_state_model_dump_json(self):
|
|
"""WorkflowState 支持 model_dump_json"""
|
|
from kb_pipeline.schemas import WorkflowState
|
|
import json
|
|
|
|
state = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
json_str = state.model_dump_json()
|
|
data = json.loads(json_str)
|
|
|
|
assert data["content_id"] == "BV1test"
|
|
assert data["platform"] == "bilibili"
|
|
assert data["source_url"] == "https://test.com"
|
|
assert data["steps"] == {}
|
|
|
|
def test_workflow_state_model_validate_from_json(self):
|
|
"""WorkflowState 可从 JSON 恢复"""
|
|
from kb_pipeline.schemas import WorkflowState
|
|
import json
|
|
|
|
original = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
json_str = original.model_dump_json()
|
|
restored = WorkflowState.model_validate_json(json_str)
|
|
|
|
assert restored.content_id == original.content_id
|
|
assert restored.platform == original.platform
|
|
assert restored.source_url == original.source_url
|
|
|
|
def test_workflow_state_with_steps_serialization(self):
|
|
"""WorkflowState steps 字典正确序列化/反序列化"""
|
|
from kb_pipeline.schemas import WorkflowState, StepRecord, StepStatus
|
|
import json
|
|
|
|
original = WorkflowState(
|
|
content_id="BV1test",
|
|
platform="bilibili",
|
|
source_url="https://test.com"
|
|
)
|
|
original.steps["fetch_metadata"] = StepRecord(
|
|
step_name="fetch_metadata",
|
|
status=StepStatus.COMPLETED
|
|
)
|
|
|
|
json_str = original.model_dump_json()
|
|
data = json.loads(json_str)
|
|
|
|
assert "fetch_metadata" in data["steps"]
|
|
assert data["steps"]["fetch_metadata"]["status"] == "completed"
|
|
|
|
restored = WorkflowState.model_validate_json(json_str)
|
|
assert "fetch_metadata" in restored.steps
|
|
assert restored.steps["fetch_metadata"].status == StepStatus.COMPLETED
|