375 lines
12 KiB
Python
375 lines
12 KiB
Python
"""Tests for generation harness runtime support."""
|
|
|
|
import pytest
|
|
from sqlalchemy import select
|
|
|
|
from app.db.models import GenerationJob, GenerationJobEvent
|
|
from app.services.adapters.storybook.primary import Storybook, StorybookPage
|
|
from app.services.adapters.text.models import StoryOutput
|
|
from app.services.generation_jobs import create_generation_job, record_generation_event
|
|
from app.services.harness.control import ExecutionControl, GenerationJobCanceledError
|
|
from app.services.harness.plans import (
|
|
build_asset_plan,
|
|
build_story_plan,
|
|
build_storybook_plan,
|
|
)
|
|
from app.services.harness.quality_gates import (
|
|
QualityGateError,
|
|
validate_story_output,
|
|
validate_storybook_output,
|
|
)
|
|
from app.services.harness.trace import TraceRecorder
|
|
from app.services.harness.types import (
|
|
ArtifactKind,
|
|
FailureCategory,
|
|
WorkflowStep,
|
|
artifact_for_event,
|
|
normalize_trace_metadata,
|
|
step_for_event,
|
|
)
|
|
|
|
|
|
def test_event_type_maps_to_standard_workflow_step():
|
|
assert step_for_event("request_accepted") == WorkflowStep.REQUEST_ACCEPTANCE
|
|
assert step_for_event("context_prepared") == WorkflowStep.CONTEXT_PREPARATION
|
|
assert step_for_event("narrative_generated") == WorkflowStep.NARRATIVE_GENERATION
|
|
assert step_for_event("story_saved") == WorkflowStep.STORY_PERSISTENCE
|
|
assert step_for_event("provider_call_succeeded") == WorkflowStep.PROVIDER_INVOCATION
|
|
assert step_for_event("quality_gate_failed") == WorkflowStep.NARRATIVE_GENERATION
|
|
assert step_for_event("cover_image_failed") == WorkflowStep.IMAGE_GENERATION
|
|
assert step_for_event("audio_succeeded") == WorkflowStep.AUDIO_GENERATION
|
|
assert step_for_event("generation_canceled") == WorkflowStep.CANCELLATION
|
|
assert step_for_event("generation_stale_failed") == WorkflowStep.STALE_RECOVERY
|
|
assert step_for_event("future_event") == WorkflowStep.UNKNOWN
|
|
|
|
|
|
def test_event_type_maps_to_standard_artifact():
|
|
assert artifact_for_event("narrative_generated") == ArtifactKind.STORY_TEXT
|
|
assert artifact_for_event("quality_gate_failed") == ArtifactKind.STORY_TEXT
|
|
assert artifact_for_event("cover_image_succeeded") == ArtifactKind.COVER_IMAGE
|
|
assert artifact_for_event("storybook_page_image_failed") == ArtifactKind.PAGE_IMAGE
|
|
assert artifact_for_event("audio_cache_hit") == ArtifactKind.AUDIO
|
|
assert artifact_for_event("postprocessing_queued") == ArtifactKind.ACHIEVEMENT_MEMORY
|
|
assert artifact_for_event("request_accepted") == ArtifactKind.NONE
|
|
|
|
|
|
def test_trace_metadata_adds_standard_fields_without_dropping_legacy_values():
|
|
metadata = normalize_trace_metadata(
|
|
"provider_call_failed",
|
|
{
|
|
"capability": "text",
|
|
"adapter": "demo",
|
|
"error": "timeout",
|
|
},
|
|
failure_category=FailureCategory.TIMEOUT,
|
|
retryable=True,
|
|
)
|
|
|
|
assert metadata["capability"] == "text"
|
|
assert metadata["adapter"] == "demo"
|
|
assert metadata["error"] == "timeout"
|
|
assert metadata["step"] == "provider_invocation"
|
|
assert metadata["artifact"] == "none"
|
|
assert metadata["failure_category"] == "timeout"
|
|
assert metadata["retryable"] is True
|
|
assert metadata["blocks_main_result"] is False
|
|
|
|
|
|
def test_trace_metadata_respects_explicit_step_and_artifact():
|
|
metadata = normalize_trace_metadata(
|
|
"narrative_generated",
|
|
{"title": "小兔子的冒险"},
|
|
step=WorkflowStep.NARRATIVE_GENERATION,
|
|
artifact=ArtifactKind.STORYBOOK_PAGES,
|
|
blocks_main_result=True,
|
|
)
|
|
|
|
assert metadata["title"] == "小兔子的冒险"
|
|
assert metadata["step"] == "narrative_generation"
|
|
assert metadata["artifact"] == "storybook_pages"
|
|
assert metadata["blocks_main_result"] is True
|
|
|
|
|
|
def test_story_plan_without_assets_snapshot():
|
|
assert build_story_plan(generate_images=False).to_snapshot() == {
|
|
"mode": "story",
|
|
"tasks": [
|
|
{
|
|
"key": "prepare_context",
|
|
"step": "context_preparation",
|
|
"artifact": "none",
|
|
"required": True,
|
|
"recoverable": False,
|
|
},
|
|
{
|
|
"key": "generate_narrative",
|
|
"step": "narrative_generation",
|
|
"artifact": "story_text",
|
|
"required": True,
|
|
"recoverable": False,
|
|
},
|
|
{
|
|
"key": "persist_story",
|
|
"step": "story_persistence",
|
|
"artifact": "story_text",
|
|
"required": True,
|
|
"recoverable": False,
|
|
},
|
|
{
|
|
"key": "queue_postprocessing",
|
|
"step": "postprocessing",
|
|
"artifact": "achievement_memory",
|
|
"required": False,
|
|
"recoverable": True,
|
|
},
|
|
{
|
|
"key": "complete_generation",
|
|
"step": "completion",
|
|
"artifact": "none",
|
|
"required": True,
|
|
"recoverable": False,
|
|
},
|
|
],
|
|
}
|
|
|
|
|
|
def test_story_plan_with_assets_marks_cover_recoverable():
|
|
plan = build_story_plan(generate_images=True).to_snapshot()
|
|
|
|
assert plan["mode"] == "story_with_assets"
|
|
assert plan["tasks"][3] == {
|
|
"key": "generate_cover_image",
|
|
"step": "image_generation",
|
|
"artifact": "cover_image",
|
|
"required": False,
|
|
"recoverable": True,
|
|
}
|
|
|
|
|
|
def test_storybook_plan_with_images_marks_storybook_images_recoverable():
|
|
plan = build_storybook_plan(generate_images=True).to_snapshot()
|
|
|
|
assert plan["mode"] == "storybook"
|
|
assert [task["key"] for task in plan["tasks"]] == [
|
|
"prepare_context",
|
|
"generate_storybook_pages",
|
|
"generate_storybook_images",
|
|
"persist_storybook",
|
|
"queue_postprocessing",
|
|
"complete_generation",
|
|
]
|
|
assert plan["tasks"][2]["artifact"] == "image"
|
|
assert plan["tasks"][2]["recoverable"] is True
|
|
|
|
|
|
def test_asset_retry_plan_deduplicates_assets():
|
|
plan = build_asset_plan(output_mode="asset_retry", assets=["image", "audio", "image"])
|
|
|
|
assert plan.to_snapshot() == {
|
|
"mode": "asset_retry",
|
|
"tasks": [
|
|
{
|
|
"key": "start_asset_retry",
|
|
"step": "asset_retry",
|
|
"artifact": "none",
|
|
"required": True,
|
|
"recoverable": False,
|
|
},
|
|
{
|
|
"key": "complete_image_asset",
|
|
"step": "image_generation",
|
|
"artifact": "image",
|
|
"required": False,
|
|
"recoverable": True,
|
|
},
|
|
{
|
|
"key": "complete_audio_asset",
|
|
"step": "audio_generation",
|
|
"artifact": "audio",
|
|
"required": False,
|
|
"recoverable": True,
|
|
},
|
|
{
|
|
"key": "complete_asset_retry",
|
|
"step": "asset_retry",
|
|
"artifact": "none",
|
|
"required": True,
|
|
"recoverable": False,
|
|
},
|
|
],
|
|
}
|
|
|
|
|
|
def test_story_quality_gate_accepts_complete_child_safe_story():
|
|
validate_story_output(
|
|
StoryOutput(
|
|
mode="generated",
|
|
title="小兔子的月光花园",
|
|
story_text="小兔子在花园里学会了和朋友轮流分享水壶。",
|
|
cover_prompt_suggestion="A gentle moonlit garden with a rabbit",
|
|
)
|
|
)
|
|
|
|
|
|
def test_story_quality_gate_rejects_missing_story_text():
|
|
output = StoryOutput(
|
|
mode="generated",
|
|
title="空白故事",
|
|
story_text="",
|
|
cover_prompt_suggestion="A cover",
|
|
)
|
|
|
|
try:
|
|
validate_story_output(output)
|
|
except QualityGateError as exc:
|
|
assert [issue.code.value for issue in exc.issues] == ["missing_story_text"]
|
|
assert exc.to_metadata()["issues"][0]["field"] == "story_text"
|
|
else:
|
|
raise AssertionError("Expected QualityGateError")
|
|
|
|
|
|
def test_story_quality_gate_rejects_obviously_unsafe_child_content():
|
|
output = StoryOutput(
|
|
mode="generated",
|
|
title="危险词测试",
|
|
story_text="这个故事包含血腥场景。",
|
|
cover_prompt_suggestion="A cover",
|
|
)
|
|
|
|
try:
|
|
validate_story_output(output)
|
|
except QualityGateError as exc:
|
|
assert [issue.code.value for issue in exc.issues] == ["unsafe_child_content"]
|
|
assert exc.to_metadata()["issues"][0]["failure_category"] == "safety_error"
|
|
else:
|
|
raise AssertionError("Expected QualityGateError")
|
|
|
|
|
|
def test_storybook_quality_gate_rejects_duplicate_page_number():
|
|
storybook = Storybook(
|
|
title="森林绘本",
|
|
main_character="小兔子",
|
|
art_style="水彩",
|
|
cover_prompt="A forest cover",
|
|
pages=[
|
|
StorybookPage(page_number=1, text="第一页。", image_prompt="page 1"),
|
|
StorybookPage(page_number=1, text="第二页。", image_prompt="page 2"),
|
|
],
|
|
)
|
|
|
|
try:
|
|
validate_storybook_output(storybook)
|
|
except QualityGateError as exc:
|
|
assert [issue.code.value for issue in exc.issues] == [
|
|
"invalid_storybook_page_number"
|
|
]
|
|
assert exc.to_metadata()["issues"][0]["field"] == "pages[1].page_number"
|
|
else:
|
|
raise AssertionError("Expected QualityGateError")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_trace_recorder_persists_standard_metadata(db_session, test_user):
|
|
job = await create_generation_job(
|
|
db_session,
|
|
user_id=test_user.id,
|
|
output_mode="story",
|
|
input_type="keywords",
|
|
request_payload={"data": "小兔子"},
|
|
)
|
|
|
|
event = await TraceRecorder(db_session).record_step(
|
|
job=job,
|
|
event_type="provider_call_failed",
|
|
status="failed",
|
|
metadata={
|
|
"capability": "text",
|
|
"adapter": "demo",
|
|
"error": "timeout",
|
|
},
|
|
failure_category=FailureCategory.TIMEOUT,
|
|
retryable=True,
|
|
)
|
|
|
|
assert event is not None
|
|
events = (
|
|
await db_session.execute(
|
|
select(GenerationJobEvent)
|
|
.where(GenerationJobEvent.job_id == job.id)
|
|
.order_by(GenerationJobEvent.id)
|
|
)
|
|
).scalars().all()
|
|
|
|
assert [item.event_type for item in events] == [
|
|
"request_accepted",
|
|
"provider_call_failed",
|
|
]
|
|
metadata = events[1].event_metadata
|
|
assert metadata["capability"] == "text"
|
|
assert metadata["adapter"] == "demo"
|
|
assert metadata["step"] == "provider_invocation"
|
|
assert metadata["artifact"] == "none"
|
|
assert metadata["failure_category"] == "timeout"
|
|
assert metadata["retryable"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_trace_recorder_ignores_missing_job(db_session):
|
|
event = await TraceRecorder(db_session).record_step(
|
|
job=None,
|
|
event_type="context_prepared",
|
|
status="succeeded",
|
|
)
|
|
|
|
assert event is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execution_control_cancels_job_at_safe_checkpoint(
|
|
db_session,
|
|
test_user,
|
|
test_story,
|
|
):
|
|
job = await create_generation_job(
|
|
db_session,
|
|
user_id=test_user.id,
|
|
output_mode="story",
|
|
input_type="keywords",
|
|
request_payload={"data": "小兔子"},
|
|
story_id=test_story.id,
|
|
)
|
|
await record_generation_event(
|
|
db_session,
|
|
job=job,
|
|
story_id=test_story.id,
|
|
event_type="cancel_requested",
|
|
status="running",
|
|
message="Cancellation requested.",
|
|
)
|
|
|
|
with pytest.raises(GenerationJobCanceledError):
|
|
await ExecutionControl(db_session).stop_if_cancel_requested(
|
|
job=job,
|
|
story=test_story,
|
|
)
|
|
|
|
refreshed_job = (
|
|
await db_session.execute(select(GenerationJob).where(GenerationJob.id == job.id))
|
|
).scalar_one()
|
|
assert refreshed_job.status == "canceled"
|
|
assert refreshed_job.current_step == "generation_canceled"
|
|
assert refreshed_job.error_message == "Generation canceled by user."
|
|
|
|
events = (
|
|
await db_session.execute(
|
|
select(GenerationJobEvent)
|
|
.where(GenerationJobEvent.job_id == job.id)
|
|
.order_by(GenerationJobEvent.id)
|
|
)
|
|
).scalars().all()
|
|
assert [item.event_type for item in events] == [
|
|
"request_accepted",
|
|
"cancel_requested",
|
|
"generation_canceled",
|
|
]
|