eryao/backend/tests/unit/test_json_finalize.py

from __future__ import annotations

import json
from typing import Any

import pytest
from pydantic import BaseModel, ConfigDict, Field

from core.agentscope.utils.json_finalize import (
    build_json_finalize_instruction,
    finalize_json_response,
)


class _Inner(BaseModel):
    model_config = ConfigDict(extra="forbid")

    year_gan_zhi: str = Field(alias="yearGanZhi")


class _Output(BaseModel):
    model_config = ConfigDict(extra="forbid", populate_by_name=True)

    ganzhi: _Inner


class _Formatter:
    async def format(self, *args: Any, **kwargs: Any) -> Any:
        del args, kwargs
        return [{"role": "user", "content": "prompt"}]


class _Response:
    def __init__(self, payload: dict[str, Any]) -> None:
        self.content = [
            {"type": "text", "text": json.dumps(payload, ensure_ascii=False)}
        ]


class _Model:
    def __init__(self, payload: dict[str, Any]) -> None:
        self._payload = payload
        self.stream = False

    async def __call__(self, *args: Any, **kwargs: Any) -> _Response:
        del args, kwargs
        return _Response(self._payload)


def test_build_instruction_has_schema() -> None:
    instruction = build_json_finalize_instruction(
        schema_json="{}",
        attempt=1,
    )
    assert "[Output Schema]" in instruction


def test_build_instruction_has_language_constraint_en() -> None:
    instruction = build_json_finalize_instruction(
        schema_json="{}",
        attempt=1,
        language="en-US",
    )
    assert "ENGLISH OUTPUT REQUIRED" in instruction
    assert "═══════════════════════════════════════════════════════════════════════════════" in instruction


def test_build_instruction_has_language_constraint_zh() -> None:
    instruction = build_json_finalize_instruction(
        schema_json="{}",
        attempt=1,
        language="zh-CN",
    )
    assert "返回 JSON。使用简体中文" in instruction


def test_build_instruction_no_language_constraint_when_none() -> None:
    instruction = build_json_finalize_instruction(
        schema_json="{}",
        attempt=1,
        language=None,
    )
    assert "ENGLISH OUTPUT REQUIRED" not in instruction
    assert "返回 JSON" not in instruction


@pytest.mark.asyncio
async def test_finalize_json_response_returns_alias_keys() -> None:
    model = _Model(payload={"ganzhi": {"yearGanZhi": "丙午"}})
    _, result = await finalize_json_response(
        model=model,
        formatter=_Formatter(),
        base_messages=[],
        output_model=_Output,
        retries=0,
    )

    assert result.model_dump(mode="json", by_alias=True) == {"ganzhi": {"yearGanZhi": "丙午"}}