AnyImage/tests/test_ai_image_operators.py

from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch
from blender_test_support import BlenderTestCase


class AiImageOperatorsTest(BlenderTestCase):
    def test_upscale_operator_defaults_to_wdn_and_sends_selected_model(self):
        operator_type = self.upscale.UpscaleImage
        self.assertIn(operator_type, self.anyimage.CLASSES)
        self.assertEqual(
            operator_type.__annotations__["model"]["default"],
            "REALESRGAN_GENERAL_WDN_X4V3",
        )
        self.assertEqual(
            operator_type.__annotations__["max_input_size"]["default"],
            2048,
        )

        self.assertIn(self.runtime.JobOperatorBase, operator_type.__mro__)
        self.assertEqual(operator_type.bl_options, {"UNDO"})
        operator = operator_type()
        operator.input_path = "input.png"
        operator.source_object_name = "Source"
        operator.model = "SAFMN_REAL_X4"
        operator.max_input_size = 1024
        context = SimpleNamespace(
            scene=SimpleNamespace(anyimage_settings=SimpleNamespace(video_frames=7))
        )
        with (
            patch.object(self.upscale, "require_environment"),
            patch.object(self.upscale, "require_model") as require_model,
            patch.object(
                self.upscale, "require_input_path", return_value=Path("input.png")
            ),
            patch.object(self.upscale, "production_device", return_value="CUDA"),
            patch.object(self.upscale, "source_start_frame", return_value=4),
        ):
            request = operator.request(context)

        require_model.assert_called_once_with("SAFMN_REAL_X4")
        self.assertEqual(
            request,
            {
                "input": "input.png",
                "model": "SAFMN_REAL_X4",
                "device": "cuda",
                "video_frames": 7,
                "start_frame": 4,
                "max_input_size": 1024,
            },
        )


    def test_upscale_operator_opens_setup_when_required_ai_is_missing(self):
        expected = {"CANCELLED"}
        with patch.object(
            self.upscale,
            "invoke_ai_setup_if_needed",
            return_value=expected,
        ) as setup:
            result = self.upscale.UpscaleImage().invoke(None, None)

        setup.assert_called_once_with()
        self.assertIs(result, expected)


    def test_upscale_operator_runs_immediately_when_ai_is_ready(self):
        operator = self.upscale.UpscaleImage()
        context = object()
        expected = {"FINISHED"}
        with (
            patch.object(
                self.upscale,
                "invoke_ai_setup_if_needed",
                return_value=None,
            ),
            patch.object(
                operator,
                "execute",
                return_value=expected,
            ) as execute,
        ):
            result = operator.invoke(context, None)

        execute.assert_called_once_with(context)
        self.assertIs(result, expected)
        self.assertNotIn("REGISTER", operator.bl_options)


    def test_remove_background_loads_a_plain_image_edit_result(self):
        source = SimpleNamespace()
        result_image = SimpleNamespace()
        result = SimpleNamespace(directory=Path("result"))
        ops = SimpleNamespace(
            ed=SimpleNamespace(undo_push=Mock(return_value={"FINISHED"}))
        )
        with (
            patch.object(self.image_data, "require_image_empty", return_value=source),
            patch.object(
                self.image_data,
                "load_image_edit_result",
                return_value=result_image,
            ) as load,
            patch.object(self.image_data, "replace_empty_image") as replace,
            patch.object(self.remove_background.bpy, "ops", ops, create=True),
        ):
            self.remove_background.apply_background_response(result, "Source")

        load.assert_called_once_with(source, Path("result"))
        replace.assert_called_once_with(source, result_image)


    def test_upscale_loads_a_plain_image_edit_result(self):
        source_image = SimpleNamespace()
        source = SimpleNamespace(data=source_image)
        result_image = SimpleNamespace()
        result = SimpleNamespace(directory=Path("result"))
        operator = self.upscale.UpscaleImage()
        operator.source_object_name = "Source"
        operator.source_image_name = "Poster.png"
        operator.model = "SAFMN_REAL_X4"
        images = SimpleNamespace(get=lambda name: source_image if name == "Poster.png" else None)
        with (
            patch.object(self.upscale, "require_image_empty", return_value=source),
            patch.object(
                self.upscale,
                "load_image_edit_result",
                return_value=result_image,
            ) as load,
            patch.object(self.upscale, "replace_empty_image") as replace,
            patch.object(
                self.upscale,
                "upscale_model_label",
                return_value="SAFMN Real x4",
            ),
            patch.object(
                self.upscale.bpy,
                "data",
                SimpleNamespace(images=images),
            ),
        ):
            message = operator.response(None, result)

        load.assert_called_once_with(
            source,
            Path("result"),
            frame_prefix="upscale.",
        )
        replace.assert_called_once_with(source, result_image)
        self.assertEqual(message, "Image is upscaled with SAFMN Real x4")