Spaces:
Running
Running
| import ast | |
| import io | |
| import json | |
| import os | |
| import tempfile | |
| import time | |
| import unittest | |
| from html import escape | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| from unittest.mock import Mock | |
| from PIL import Image | |
| from image_inversion_job_api import ( | |
| ImageInversionJobRequest, | |
| ImageInversionJobSettings, | |
| decode_image_bytes, | |
| ) | |
| class AppV1IntegrationTests(unittest.TestCase): | |
| def setUp(self): | |
| """隔离加载应用定义，为集成测试准备六类模型结果。""" | |
| # 只加载本次集成涉及的实际业务定义，避免测试时启动 Gradio 或下载权重。 | |
| app_path = Path(__file__).resolve().parents[1] / "app.py" | |
| self.tree = ast.parse(app_path.read_text(encoding="utf-8")) | |
| definition_names = { | |
| "Tagger", "ImageValidationError", "_format_size", "validate_and_open_image", | |
| "format_tags_html", "generate_summary_text_content", "execute_image_inversion_job", | |
| } | |
| constant_names = { | |
| "MAX_UPLOAD_BYTES", "MAX_IMAGE_SIDE", "MAX_IMAGE_PIXELS", | |
| "MAX_DECOMPRESSED_BYTES", "ALLOWED_IMAGE_FORMATS", "ASSETS_REPO_ID", | |
| "ASSETS_REVISION", | |
| } | |
| definitions = [ | |
| node for node in ast.walk(self.tree) | |
| if ( | |
| isinstance(node, (ast.ClassDef, ast.FunctionDef)) | |
| and node.name in definition_names | |
| ) or ( | |
| isinstance(node, ast.Assign) | |
| and any(isinstance(target, ast.Name) and target.id in constant_names for target in node.targets) | |
| ) | |
| ] | |
| self.scope = { | |
| "Image": Image, "ImageInversionJobRequest": ImageInversionJobRequest, | |
| "os": os, "time": time, "json": json, "escape": escape, | |
| } | |
| exec(compile(ast.Module(body=definitions, type_ignores=[]), str(app_path), "exec"), self.scope) | |
| self.tagger = self.scope["Tagger"].__new__(self.scope["Tagger"]) | |
| self.tagger.device = "cpu" | |
| self.handler_output = { | |
| "feature_scores": {"blue_eyes": 0.71, "long_hair": 0.98}, | |
| "character_scores": {"test_character": 0.84}, | |
| "ip_scores": {"test_series": 0.76}, | |
| "style_scores": {"anime": 0.89}, | |
| "meta_scores": {"highres": 0.91}, | |
| "rating_scores": {"general": 0.99}, | |
| "_params": {"general_threshold": 0.17, "character_threshold": 0.27}, | |
| "_timings": {"inference_s": 0.125}, | |
| } | |
| self.tagger.handler = Mock(return_value=self.handler_output) | |
| self.scope["tagger_instance"] = self.tagger | |
| def test_predict_preserves_all_categories_and_real_ip_confidence(self): | |
| """验证完整类别元数据和真实 IP 置信度。""" | |
| image = Image.new("RGBA", (32, 32), (10, 20, 30, 0)) | |
| result, translation_order, metadata = self.tagger.predict(image) | |
| self.assertEqual(list(result), ["general", "characters", "ips"]) | |
| self.assertEqual(result["general"], {"long hair": 0.98, "blue eyes": 0.71}) | |
| self.assertEqual(translation_order["general"], ["long hair", "blue eyes"]) | |
| self.assertEqual(result["ips"], {"test series": 0.76}) | |
| self.assertEqual(metadata["tags"], { | |
| "general": self.handler_output["feature_scores"], | |
| "character": self.handler_output["character_scores"], | |
| "copyright": self.handler_output["ip_scores"], | |
| "style": self.handler_output["style_scores"], | |
| "meta": self.handler_output["meta_scores"], | |
| "rating": self.handler_output["rating_scores"], | |
| }) | |
| call_data = self.tagger.handler.call_args.args[0] | |
| self.assertIs(call_data["inputs"], image) | |
| self.assertTrue(call_data["parameters"]["include_scores"]) | |
| self.assertEqual(call_data["parameters"]["general_threshold"], 0.17) | |
| self.assertEqual(call_data["parameters"]["character_threshold"], 0.27) | |
| def test_job_keeps_six_fields_and_translation_alignment(self): | |
| """验证任务接口保持六个字段且翻译不串位。""" | |
| self.scope["translate_texts"] = Mock(return_value=["长发", "蓝眼睛", "测试角色", "测试作品"]) | |
| payload = ImageInversionJobRequest(input_image_url="https://example.com/image.png") | |
| result = self.scope["execute_image_inversion_job"](payload, Image.new("RGB", (32, 32)), "job-id") | |
| self.assertEqual(set(result), { | |
| "status_markdown", "general_tags_html", "character_tags_html", | |
| "ip_tags_html", "summary_text", "metadata", | |
| }) | |
| self.scope["translate_texts"].assert_called_once_with( | |
| ["long hair", "blue eyes", "test character", "test series"], | |
| src_lang="auto", tgt_lang="zh", | |
| ) | |
| self.assertIn("长发", result["general_tags_html"]) | |
| self.assertIn("测试角色", result["character_tags_html"]) | |
| self.assertIn("测试作品", result["ip_tags_html"]) | |
| self.assertIn("0.760", result["ip_tags_html"]) | |
| self.assertIn("test series/*测试作品*/", result["summary_text"]) | |
| self.assertEqual(len(result["metadata"]["tags"]), 6) | |
| # 整个翻译调用异常也不能影响标签和新增元数据的交付。 | |
| self.scope["translate_texts"].side_effect = RuntimeError("translation unavailable") | |
| untranslated = self.scope["execute_image_inversion_job"](payload, Image.new("RGB", (32, 32)), "job-id") | |
| self.assertEqual(set(untranslated), set(result)) | |
| self.assertIn("test series", untranslated["summary_text"]) | |
| self.assertNotIn("测试作品", untranslated["ip_tags_html"]) | |
| self.assertEqual(untranslated["metadata"]["tags"], result["metadata"]["tags"]) | |
| def test_ui_and_api_defaults_match_v1_thresholds(self): | |
| """验证页面和任务接口的默认阈值与 v1.0 推荐值一致。""" | |
| slider_nodes = [ | |
| node for node in ast.walk(self.tree) | |
| if isinstance(node, ast.Assign) | |
| and any(isinstance(target, ast.Name) and target.id in {"gen_slider", "char_slider"} for target in node.targets) | |
| ] | |
| self.scope["gr"] = SimpleNamespace(Slider=lambda *args, **kwargs: kwargs) | |
| exec(compile(ast.Module(body=slider_nodes, type_ignores=[]), "app-sliders", "exec"), self.scope) | |
| payload = ImageInversionJobRequest(input_image_url="https://example.com/image.png") | |
| self.assertEqual(self.scope["gen_slider"]["value"], 0.17) | |
| self.assertEqual(self.scope["char_slider"]["value"], 0.27) | |
| self.assertEqual(payload.general_threshold, 0.17) | |
| self.assertEqual(payload.character_threshold, 0.27) | |
| def test_upload_and_job_decode_preserve_transparency(self): | |
| """验证上传和远程解码均将透明背景留给官方处理器。""" | |
| settings = ImageInversionJobSettings(api_key="x" * 32, allowed_hosts=frozenset({"example.com"})) | |
| for mode in ("RGBA", "P"): | |
| with self.subTest(mode=mode): | |
| image = Image.new("RGBA", (32, 32), (10, 20, 30, 0)) | |
| if mode == "P": | |
| image = image.convert("P") | |
| buffer = io.BytesIO() | |
| image.save(buffer, format="PNG") | |
| decoded = decode_image_bytes(buffer.getvalue(), "image/png", settings) | |
| self.assertEqual(decoded.convert("RGBA").getpixel((0, 0))[3], 0) | |
| # 仅创建并显式删除当前测试的单个文件，不创建需递归清理的目录。 | |
| with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as source: | |
| source.write(buffer.getvalue()) | |
| source_path = Path(source.name).resolve() | |
| try: | |
| uploaded = self.scope["validate_and_open_image"](str(source_path)) | |
| self.assertEqual(uploaded.convert("RGBA").getpixel((0, 0))[3], 0) | |
| finally: | |
| source_path.unlink() | |
| if __name__ == "__main__": | |
| unittest.main() | |