Blame for tests/test_app_v1.py

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