chendandan_homework/tests/test_supervisor_nodes.py

from __future__ import annotations

from app.graph.supervisor_nodes import (
    make_supervisor_node,
)
from app.schemas.review import (
    ReviewIssue,
    TripReviewResult,
)
from app.schemas.validation import (
    BudgetEstimate,
    PlanValidationResult,
    ValidationIssue,
)


def build_validation(
    *,
    valid: bool,
) -> PlanValidationResult:
    issues = []

    if not valid:
        issues = [
            ValidationIssue(
                code="TEST_ERROR",
                severity="error",
                message="测试校验错误。",
            )
        ]

    return PlanValidationResult(
        is_valid=valid,
        error_count=0 if valid else 1,
        warning_count=0,
        issues=issues,
        budget=BudgetEstimate(
            currency="CNY",
            known_total=3000,
            user_budget=8000,
            remaining_known_budget=5000,
            status="within_known_budget",
            coverage_note="测试预算。",
        ),
        revision_feedback=(
            []
            if valid
            else ["修复测试校验错误。"]
        ),
    )


def build_review(
    *,
    approved: bool,
) -> TripReviewResult:
    if approved:
        return TripReviewResult(
            approved=True,
            summary="方案通过。",
        )

    return TripReviewResult(
        approved=False,
        summary="方案未通过。",
        issues=[
            ReviewIssue(
                category="pace",
                severity="blocking",
                message="行程过于紧凑。",
                suggestion="减少每日景点数量。",
            )
        ],
        revision_feedback=[
            "减少每日景点数量。",
        ],
    )


def test_supervisor_approves_valid_plan() -> None:
    node = make_supervisor_node(
        max_planning_attempts=3
    )

    result = node(
        {
            "planning_attempts": 1,
            "plan_validation_result": (
                build_validation(valid=True)
            ),
            "trip_review_result": (
                build_review(approved=True)
            ),
            "review_passed": True,
            "errors": [],
        }
    )

    assert (
        result["supervisor_decision"]
        == "finalize"
    )

    assert (
        result["workflow_status"]
        == "approved"
    )


def test_supervisor_requests_replan() -> None:
    node = make_supervisor_node(
        max_planning_attempts=3
    )

    result = node(
        {
            "planning_attempts": 1,
            "plan_validation_result": (
                build_validation(valid=False)
            ),
            "trip_review_result": (
                build_review(approved=False)
            ),
            "review_passed": False,
            "revision_feedback": [],
            "errors": [],
        }
    )

    assert (
        result["supervisor_decision"]
        == "replan"
    )

    assert result["revision_feedback"]


def test_supervisor_stops_after_max_attempts(
) -> None:
    node = make_supervisor_node(
        max_planning_attempts=3
    )

    result = node(
        {
            "planning_attempts": 3,
            "plan_validation_result": (
                build_validation(valid=False)
            ),
            "trip_review_result": (
                build_review(approved=False)
            ),
            "review_passed": False,
            "revision_feedback": [],
            "errors": [],
        }
    )

    assert (
        result["supervisor_decision"]
        == "finalize_with_risks"
    )

    assert (
        result["workflow_status"]
        == "completed_with_risks"
    )


def test_supervisor_terminates_system_error(
) -> None:
    node = make_supervisor_node(
        max_planning_attempts=3
    )

    result = node(
        {
            "planning_attempts": 1,
            "errors": [
                "MCP连接失败。",
            ],
        }
    )

    assert (
        result["supervisor_decision"]
        == "terminate"
    )

    assert (
        result["workflow_status"]
        == "failed"
    )