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()