eryao/backend/tests/unit/test_runtime_models_worker_output.py

from __future__ import annotations

import pytest
from pydantic import ValidationError

from schemas.agent.runtime_models import (
    FollowUpOutput,
    WorkerAgentOutputLite,
    resolve_worker_output_model,
)


def test_worker_output_lite_rejects_divination_derived_from_llm() -> None:
    payload = {
        "status": "success",
        "sign_level": "中上签",
        "conclusion": ["结论"],
        "focus_points": ["重点"],
        "advice": ["建议"],
        "keywords": ["关键词一", "关键词二", "关键词三"],
        "answer": "最终答案",
        "divination_derived": {},
    }

    with pytest.raises(ValidationError):
        WorkerAgentOutputLite.model_validate(payload)


def test_follow_up_output_accepts_minimal_schema() -> None:
    payload = {
        "status": "success",
        "answer": "追问回答",
    }

    parsed = FollowUpOutput.model_validate(payload)
    assert parsed.answer == "追问回答"


def test_follow_up_output_rejects_chat_only_fields() -> None:
    payload = {
        "status": "success",
        "answer": "追问回答",
        "sign_level": "中上签",
    }

    with pytest.raises(ValidationError):
        FollowUpOutput.model_validate(payload)


def test_resolve_worker_output_model_uses_runtime_mode() -> None:
    assert resolve_worker_output_model(runtime_mode="chat") is WorkerAgentOutputLite
    assert resolve_worker_output_model(runtime_mode="follow_up") is FollowUpOutput