tests/test_character_spike.py

from __future__ import annotations

import json
import unittest
from pathlib import Path

from scripts.run_character_generation_spike import _evaluate_gate
from src.character_spike.assets import create_mock_asset_package
from src.character_spike.schema import EXPRESSIONS, MODEL_CANDIDATES, candidate_by_id
from src.character_spike.tavern_import import convert_tavern_card, load_tavern_json, write_character_draft


PROJECT_ROOT = Path(__file__).resolve().parents[1]
TMP_ROOT = PROJECT_ROOT / ".tmp" / "tests"


class CharacterSpikeTests(unittest.TestCase):
    def test_candidate_registry_tracks_required_models(self) -> None:
        ids = {candidate.id for candidate in MODEL_CANDIDATES}
        self.assertIn("flux_schnell", ids)
        self.assertIn("qwen_image", ids)
        self.assertIn("qwen_image_edit", ids)
        self.assertIn("qwen_controlnet_union", ids)
        self.assertIn("instantid_sdxl", ids)
        self.assertFalse(candidate_by_id("instantid_sdxl").implemented)

    def test_tavern_json_import_maps_fields_to_draft_package(self) -> None:
        card = {
            "spec": "chara_card_v2",
            "spec_version": "2.0",
            "data": {
                "name": "测试角色",
                "description": "一名用于导入验证的角色。\n第二行会被保留。",
                "personality": "冷静、敏锐，温柔",
                "scenario": "用户正在测试角色卡导入。",
                "first_mes": "你好，我已经导入完成。",
                "alternate_greetings": ["备用开场 A", "备用开场 B"],
                "creator_notes": "只用于测试。",
                "tags": ["测试", "导入"],
                "character_book": {"entries": [{"keys": ["地点"], "content": "测试空间"}]},
            },
        }

        package = convert_tavern_card(card)

        self.assertEqual(package["display_name"], "测试角色")
        self.assertEqual(package["first_mes"], "你好，我已经导入完成。")
        self.assertEqual(package["alternate_greetings"], ["备用开场 A", "备用开场 B"])
        self.assertEqual(package["profile"]["core_traits"][:3], ["冷静", "敏锐", "温柔"])
        self.assertEqual(package["metadata"]["spec"], "chara_card_v2")
        self.assertEqual(package["metadata"]["character_book"]["entries"][0]["content"], "测试空间")

    def test_load_and_write_tavern_draft(self) -> None:
        root = TMP_ROOT / "tavern_draft"
        root.mkdir(parents=True, exist_ok=True)
        source = root / "card.json"
        source.write_text(json.dumps({"name": "Plain Card", "first_mes": "Hi"}), encoding="utf-8")

        package = convert_tavern_card(load_tavern_json(source), forced_id="plain_card")
        path = write_character_draft(package, root / "characters")

        self.assertTrue(path.exists())
        parsed = json.loads(path.read_text(encoding="utf-8"))
        self.assertEqual(parsed["id"], "plain_card")
        self.assertEqual(parsed["first_mes"], "Hi")

    def test_mock_asset_package_creates_manifest_grid_and_assets(self) -> None:
        output_root = TMP_ROOT / "mock_assets"
        output_root.mkdir(parents=True, exist_ok=True)
        manifest = create_mock_asset_package(
            character_id="spike_test",
            display_name="星核",
            output_root=output_root,
            seed=7,
        )
        run_dir = Path(manifest["paths"]["run_dir"])

        self.assertEqual(len(manifest["assets"]), len(EXPRESSIONS))
        self.assertTrue((run_dir / "generated" / "manifest.json").exists())
        self.assertTrue((run_dir / "generated" / "report.md").exists())
        self.assertTrue((run_dir / "generated" / "asset_grid.png").exists())
        for expression in EXPRESSIONS:
            self.assertTrue((run_dir / "assets" / "characters" / "spike_test" / f"{expression}.png").exists())

    def test_gate_requires_timing_and_usable_assets(self) -> None:
        manifest = {
            "assets": [{"usable": True} for _ in range(6)],
            "model_results": [
                {
                    "benchmark_case": "warm_four",
                    "status": "ok",
                    "loaded_before": True,
                    "duration_seconds": 55.0,
                },
                {
                    "benchmark_case": "eight_expression_total",
                    "status": "ok",
                    "duration_seconds": 170.0,
                },
            ],
        }
        gate = _evaluate_gate(manifest)
        self.assertTrue(gate["ready_for_character_workshop"])

        manifest["model_results"][1]["duration_seconds"] = 240.0
        gate = _evaluate_gate(manifest)
        self.assertFalse(gate["ready_for_character_workshop"])
        self.assertIn("导入角色卡", gate["fallback_recommendation"])


if __name__ == "__main__":
    unittest.main()