Blame for tests/test_handler.py

"""不下载模型权重，验证 v1.0 推理适配层与官方 pipeline 的调用契约。"""

import base64
import importlib.util
import io
import os
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, Mock, patch

from PIL import Image


class EndpointHandlerTest(unittest.TestCase):
    """覆盖标签映射、真实分数、预处理委托和本地加载边界。"""

    @classmethod
    def setUpClass(cls):
        """隔离导入推理依赖，让适配层测试可在没有 GPU 和权重时运行。"""
        spec = importlib.util.spec_from_file_location(
            "pixai_handler_test", Path(__file__).resolve().parents[1] / "handler.py"
        )
        cls.module = importlib.util.module_from_spec(spec)
        cls.torch = SimpleNamespace(
            cuda=SimpleNamespace(is_available=Mock(return_value=False)),
            zeros=Mock(),
        )
        with patch.dict(
            sys.modules,
            {"torch": cls.torch, "transformers": SimpleNamespace(pipeline=Mock())},
        ):
            spec.loader.exec_module(cls.module)

    def setUp(self):
        """构造各类别独立的分数，避免将版权或风格误当成通用标签。"""
        self.results = {
            "general": {"blue_hair": 0.7, "1girl": 0.99, "hat": 0.2},
            "character": {"character_b": 0.3, "character_a": 0.9},
            "copyright": {"original": 0.64},
            "style": {"style_a": 0.6},
            "meta": {"highres": 0.8},
            "rating": {"rating:g": 0.88},
        }
        self.pipe = Mock(return_value={"results": self.results})
        self.factory = Mock(return_value=self.pipe)
        self.torch.cuda.is_available.return_value = False
        self.torch.zeros.reset_mock()
        self.torch.zeros.return_value.to.side_effect = None
        self.model_path = Path(__file__).resolve().parents[1]
        with patch.object(Path, "is_file", return_value=True), patch.object(
            self.module, "pipeline", self.factory
        ):
            self.handler = self.module.EndpointHandler(str(self.model_path))

    def test_local_snapshot_and_official_processor(self):
        """本地四文件快照与官方自定义任务共同决定模型和预处理版本。"""
        self.factory.assert_called_once_with(
            task="tagger",
            model=str(self.model_path),
            image_processor=str(self.model_path),
            trust_remote_code=True,
            token=False,
            model_kwargs={"local_files_only": True},
            device="cpu",
        )
        image = Image.new("RGBA", (80, 40), (255, 0, 0, 0))
        data = {"inputs": image, "parameters": {"include_scores": True}}
        result = self.handler(data)
        self.pipe.assert_called_once_with(
            image, threshold={"general": 0.17, "character": 0.27}
        )
        self.assertEqual(image.mode, "RGBA")
        self.assertEqual(image.size, (80, 40))
        self.assertIs(data["inputs"], image)
        self.assertEqual(data["parameters"], {"include_scores": True})
        self.assertEqual(result["feature"], ["1girl", "blue_hair", "hat"])
        self.assertEqual(result["feature_scores"], self.results["general"])
        for category in ("character", "style", "meta", "rating"):
            self.assertEqual(result[f"{category}_scores"], self.results[category])
        self.assertEqual(result["ip"], ["original"])
        self.assertEqual(result["ip_scores"], {"original": 0.64})
        self.assertEqual(result["_params"]["general_threshold"], 0.17)
        self.assertEqual(set(result["_timings"]), {"fetch_s", "inference_s", "post_process_s", "total_s"})

    def test_custom_thresholds_and_optional_scores(self):
        """阈值交由官方流程筛选，关闭分数时不泄漏分数字段。"""
        image = Image.new("RGB", (32, 32))
        result = self.handler(
            {"inputs": image, "parameters": {"general_threshold": 0.4, "character_threshold": 0.6}}
        )
        self.pipe.assert_called_once_with(image, threshold={"general": 0.4, "character": 0.6})
        self.assertFalse(any(name.endswith("_scores") for name in result))

    def test_topk_ranks_categories_independently(self):
        """Top-K 不先经过默认阈值，且保留版权等其他分类的原始分数。"""
        image = Image.new("RGB", (32, 32))
        result = self.handler(
            {
                "inputs": image,
                "parameters": {"mode": "topk", "topk_general": 2, "topk_character": 1, "include_scores": True},
            }
        )
        self.pipe.assert_called_once_with(image, threshold={"general": -1.0, "character": -1.0})
        self.assertEqual(result["feature"], ["1girl", "blue_hair"])
        self.assertEqual(result["character"], ["character_a"])
        self.assertEqual(result["character_scores"], {"character_a": 0.9})
        self.assertEqual(result["ip_scores"], {"original": 0.64})
        zero_result = self.handler(
            {"inputs": image, "parameters": {"mode": "topk", "topk_general": 0, "topk_character": 0}}
        )
        self.assertEqual(zero_result["feature"], [])
        self.assertEqual(zero_result["character"], [])

    def test_base64_and_remote_images_are_decoded_before_inference(self):
        """URL 图片在响应关闭前解码，base64 和透明图片也交由官方处理器。"""
        image_bytes = io.BytesIO()
        Image.new("RGBA", (80, 40), (0, 255, 0, 0)).save(image_bytes, format="PNG")
        image_content = image_bytes.getvalue()
        self.handler({"inputs": {"image": base64.b64encode(image_content).decode()}})
        decoded_image = self.pipe.call_args.args[0]
        self.assertEqual((decoded_image.mode, decoded_image.size), ("RGBA", (80, 40)))

        response = MagicMock()
        response.raw = io.BytesIO(image_content)
        response.__enter__.return_value = response
        response.__exit__.side_effect = lambda *_: response.raw.close()
        with patch.object(self.module.requests, "get", return_value=response) as get:
            self.handler({"inputs": {"url": "https://example.org/image.png"}})
        get.assert_called_once_with("https://example.org/image.png", stream=True, timeout=5.0)
        response.raise_for_status.assert_called_once_with()
        self.assertTrue(response.raw.closed)
        self.assertEqual(self.pipe.call_args.args[0].getpixel((0, 0)), (0, 255, 0, 0))

    def test_incomplete_snapshot_fails_before_model_loading(self):
        """缺失新版本预处理器时直接报错，避免加载旧版或远程补齐混用。"""
        with patch.object(
            Path, "is_file", autospec=True, side_effect=lambda path: path.name != "preprocessor_config.json"
        ), patch.object(self.module, "pipeline") as factory:
            with self.assertRaisesRegex(FileNotFoundError, "preprocessor_config.json"):
                self.module.EndpointHandler(str(self.model_path))
        factory.assert_not_called()

    def test_cuda_probe_failure_uses_cpu(self):
        """存在 CUDA 接口但驱动不能工作时，仍能使用 CPU 加载。"""
        self.torch.cuda.is_available.return_value = True
        self.torch.zeros.return_value.to.side_effect = RuntimeError("driver unavailable")
        with patch.dict(os.environ, {"FORCE_CPU": "0"}), patch.object(
            Path, "is_file", return_value=True
        ), patch.object(self.module, "pipeline", self.factory):
            handler = self.module.EndpointHandler(str(self.model_path))
        self.assertEqual(handler.device, "cpu")
        self.assertEqual(self.factory.call_args.kwargs["device"], "cpu")


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