tests/test_character_workshop.py

from __future__ import annotations

import io
import json
import os
import unittest
import uuid
from pathlib import Path

from PIL import Image, ImageDraw

import src.character_workshop as workshop
from src.character_registry import load_user_character_packages
from src.model_status import ModelStatus, statuses_markdown


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


class CharacterWorkshopTests(unittest.TestCase):
    def setUp(self) -> None:
        self.test_root = TMP_ROOT / f"{self._testMethodName}_{uuid.uuid4().hex}"
        self.test_root.mkdir(parents=True, exist_ok=True)
        self.old_roots = (
            workshop.WORKSHOP_ROOT,
            workshop.INSTALLED_CHARACTER_ROOT,
            workshop.INSTALLED_BACKGROUND_ROOT,
            workshop.INSTALLED_CHARACTER_JSON_ROOT,
        )
        self.old_data_root = os.environ.get("VC_DATA_ROOT")
        workshop.WORKSHOP_ROOT = self.test_root / "generated"
        workshop.INSTALLED_CHARACTER_ROOT = self.test_root / "installed" / "assets" / "characters"
        workshop.INSTALLED_BACKGROUND_ROOT = self.test_root / "installed" / "assets" / "backgrounds"
        workshop.INSTALLED_CHARACTER_JSON_ROOT = self.test_root / "installed" / "characters"
        os.environ["VC_DATA_ROOT"] = str(self.test_root / "data")
        workshop.set_remote_probe_for_tests(self._fake_probe)
        self.probe_calls = []

    def tearDown(self) -> None:
        (
            workshop.WORKSHOP_ROOT,
            workshop.INSTALLED_CHARACTER_ROOT,
            workshop.INSTALLED_BACKGROUND_ROOT,
            workshop.INSTALLED_CHARACTER_JSON_ROOT,
        ) = self.old_roots
        if self.old_data_root is None:
            os.environ.pop("VC_DATA_ROOT", None)
        else:
            os.environ["VC_DATA_ROOT"] = self.old_data_root
        workshop.set_remote_probe_for_tests(None)

    def test_form_to_installable_character_package(self) -> None:
        draft = workshop.create_draft_from_form(
            display_name="工坊角色",
            description="银白短发的原创角色。",
            personality="冷静、温柔",
            scenario="通讯端测试。",
            first_mes="我在。",
            tags="测试,工坊",
        )
        state = workshop.create_initial_state(draft)
        state = workshop.generate_main_candidates(state)
        self.assertEqual(len(state["main_candidates"]), 4)

        state = workshop.select_main_candidate(state, 2)
        self.assertEqual(state["selected_candidate_index"], 2)

        state = workshop.generate_expression_pack(state)
        self.assertEqual(set(state["expression_assets"]), set(workshop.EXPRESSIONS))

        state = workshop.generate_background(state)
        self.assertTrue(Path(state["background_asset"]).exists())

        state = workshop.matte_and_package_assets(state, mode="fallback")
        package_dir = Path(state["package_dir"])
        self.assertTrue((package_dir / "generated" / "asset_grid.png").exists())
        self.assertTrue((package_dir / "generated" / "stage_smoke.html").exists())

        state = workshop.install_character_package(state)
        character_id = state["installed_character_id"]
        self.assertTrue((workshop.INSTALLED_CHARACTER_JSON_ROOT / f"{character_id}.json").exists())
        self.assertTrue((workshop.INSTALLED_BACKGROUND_ROOT / f"{character_id}_background.png").exists())
        for expression in workshop.EXPRESSIONS:
            self.assertTrue((workshop.INSTALLED_CHARACTER_ROOT / character_id / f"{expression}.png").exists())

    def test_user_character_loader_reads_installed_json(self) -> None:
        root = TMP_ROOT / "registry"
        root.mkdir(parents=True, exist_ok=True)
        path = root / "loaded_character.json"
        path.write_text(json.dumps({"id": "loaded_character", "display_name": "已加载角色"}), encoding="utf-8")

        packages = load_user_character_packages(root)

        self.assertEqual(packages["loaded_character"]["display_name"], "已加载角色")

    def test_image_generation_status_markdown(self) -> None:
        html = statuses_markdown(
            [
                ModelStatus(
                    "image_generation",
                    "sleeping",
                    "已休眠",
                    url="modal_apps/modal_character_spike.py",
                    message="Modal 图像生成服务可能已休眠或正在冷启动，请等待容器启动和模型载入后重试。",
                )
            ]
        )

        self.assertIn("Image Generation", html)
        self.assertIn("等待容器启动", html)

    def test_logged_in_runs_are_user_scoped_and_loadable(self) -> None:
        alice = workshop.get_current_user(_Profile("alice", "Alice"))
        bob = workshop.get_current_user(_Profile("bob", "Bob"))

        alice_state = workshop.create_initial_state(
            workshop.create_draft_from_form(
                display_name="Alice 角色",
                description="测试。",
                personality="冷静",
                scenario="测试。",
                first_mes="你好。",
            ),
            user=alice,
        )
        bob_state = workshop.create_initial_state(
            workshop.create_draft_from_form(
                display_name="Bob 角色",
                description="测试。",
                personality="冷静",
                scenario="测试。",
                first_mes="你好。",
            ),
            user=bob,
        )

        alice_runs = workshop.list_user_workshop_runs(alice)
        bob_runs = workshop.list_user_workshop_runs(bob)

        self.assertEqual(len(alice_runs), 1)
        self.assertEqual(len(bob_runs), 1)
        self.assertIn("Alice 角色", alice_runs[0][0])
        self.assertIn("Bob 角色", bob_runs[0][0])
        self.assertNotEqual(alice_state["run_dir"], bob_state["run_dir"])
        loaded = workshop.load_workshop_run(alice_runs[0][1], user=alice)
        self.assertEqual(loaded["character_id"], alice_state["character_id"])
        with self.assertRaises(ValueError):
            workshop.load_workshop_run(alice_runs[0][1], user=bob)

    def test_resume_main_candidates_from_manifest(self) -> None:
        user = workshop.get_current_user(_Profile("alice", "Alice"))
        state = workshop.create_initial_state(
            workshop.create_draft_from_form(
                display_name="恢复角色",
                description="测试。",
                personality="冷静",
                scenario="测试。",
                first_mes="你好。",
            ),
            user=user,
        )
        state = workshop.generate_main_candidates(state)

        loaded = workshop.load_workshop_run(state["run_dir"], user=user)

        self.assertEqual(len(loaded["main_candidates"]), 4)
        self.assertEqual(loaded["selected_candidate_index"], 0)

    def test_partial_expression_pack_only_generates_missing_slots(self) -> None:
        user = workshop.get_current_user(_Profile("alice", "Alice"))
        state = workshop.create_initial_state(
            workshop.create_draft_from_form(
                display_name="续跑角色",
                description="测试。",
                personality="冷静",
                scenario="测试。",
                first_mes="你好。",
            ),
            user=user,
        )
        run_dir = Path(state["run_dir"])
        idle_dir = run_dir / "expressions_raw" / "idle"
        idle_dir.mkdir(parents=True, exist_ok=True)
        idle_path = idle_dir / "00.png"
        idle_path.write_bytes(_png_bytes(0))
        manifest = json.loads((run_dir / "manifest.json").read_text(encoding="utf-8"))
        manifest["expression_assets_raw"] = {"idle": str(idle_path.relative_to(run_dir))}
        (run_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False), encoding="utf-8")

        self.probe_calls.clear()
        state = workshop.load_workshop_run(run_dir, user=user)
        state = workshop.generate_expression_pack(state)

        self.assertEqual(set(state["expression_assets"]), set(workshop.EXPRESSIONS))
        self.assertEqual(len(self.probe_calls), len(workshop.EXPRESSIONS) - 1)

    def test_stats_event_summary(self) -> None:
        user = workshop.get_current_user(_Profile("alice", "Alice"))

        workshop.record_workshop_event(user, "generate_main_candidates", {"stage": "main_candidates", "success": True, "duration_seconds": 1.2, "character_id": "x"})
        workshop.record_workshop_event(user, "generate_expression_pack", {"stage": "assets_ready", "success": False, "failure_reason": "boom", "character_id": "x"})
        summary = workshop.summarize_workshop_stats()

        self.assertEqual(summary["events"], 2)
        self.assertEqual(summary["users"], 1)
        self.assertEqual(summary["failures_by_stage"]["assets_ready"], 1)

    def _fake_probe(self, **kwargs):
        self.probe_calls.append(dict(kwargs))
        batch_size = int(kwargs.get("batch_size", 1))
        images = [_png_bytes(index) for index in range(batch_size)]
        return {
            "candidate_id": kwargs.get("candidate_id", "qwen_image"),
            "status": "ok",
            "image_count": batch_size,
            "duration_seconds": 0.01,
            "seed": kwargs.get("seed"),
            "images": images,
        }


def _png_bytes(index: int) -> bytes:
    image = Image.new("RGB", (256, 384), (238, 246, 248))
    draw = ImageDraw.Draw(image)
    accent = [(45, 212, 191), (96, 165, 250), (244, 114, 182), (250, 204, 21)][index % 4]
    draw.ellipse((76, 36, 180, 150), fill=(245, 222, 214))
    draw.rectangle((84, 150, 172, 330), fill=(35, 45, 60))
    draw.rectangle((112, 166, 144, 310), fill=accent)
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return buffer.getvalue()


class _Profile:
    def __init__(self, username: str, name: str):
        self.username = username
        self.name = name


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