diff --git a/tests/test_pipeline/test_schemas.py b/tests/test_pipeline/test_schemas.py index 8d2c1a8..0e200f5 100644 --- a/tests/test_pipeline/test_schemas.py +++ b/tests/test_pipeline/test_schemas.py @@ -215,6 +215,157 @@ class TestWorkflowState: 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 验证"""