ComfyUI/tests-unit/execution_test/preview_method_override_test.py

"""
Unit tests for Queue-specific Preview Method Override feature.

Tests the preview method override functionality:
- LatentPreviewMethod.from_string() method
- set_preview_method() function in latent_preview.py
- default_preview_method variable
- Integration with args.preview_method
"""
import pytest
from comfy.cli_args import args, LatentPreviewMethod
from latent_preview import set_preview_method, default_preview_method


class TestLatentPreviewMethodFromString:
    """Test LatentPreviewMethod.from_string() classmethod."""

    @pytest.mark.parametrize("value,expected", [
        ("auto", LatentPreviewMethod.Auto),
        ("latent2rgb", LatentPreviewMethod.Latent2RGB),
        ("taesd", LatentPreviewMethod.TAESD),
        ("none", LatentPreviewMethod.NoPreviews),
    ])
    def test_valid_values_return_enum(self, value, expected):
        """Valid string values should return corresponding enum."""
        assert LatentPreviewMethod.from_string(value) == expected

    @pytest.mark.parametrize("invalid", [
        "invalid",
        "TAESD",      # Case sensitive
        "AUTO",       # Case sensitive
        "Latent2RGB", # Case sensitive
        "latent",
        "",
        "default",    # default is special, not a method
    ])
    def test_invalid_values_return_none(self, invalid):
        """Invalid string values should return None."""
        assert LatentPreviewMethod.from_string(invalid) is None


class TestLatentPreviewMethodEnumValues:
    """Test LatentPreviewMethod enum has expected values."""

    def test_enum_values(self):
        """Verify enum values match expected strings."""
        assert LatentPreviewMethod.NoPreviews.value == "none"
        assert LatentPreviewMethod.Auto.value == "auto"
        assert LatentPreviewMethod.Latent2RGB.value == "latent2rgb"
        assert LatentPreviewMethod.TAESD.value == "taesd"

    def test_enum_count(self):
        """Verify exactly 4 preview methods exist."""
        assert len(LatentPreviewMethod) == 4


class TestSetPreviewMethod:
    """Test set_preview_method() function from latent_preview.py."""

    def setup_method(self):
        """Store original value before each test."""
        self.original = args.preview_method

    def teardown_method(self):
        """Restore original value after each test."""
        args.preview_method = self.original

    def test_override_with_taesd(self):
        """'taesd' should set args.preview_method to TAESD."""
        set_preview_method("taesd")
        assert args.preview_method == LatentPreviewMethod.TAESD

    def test_override_with_latent2rgb(self):
        """'latent2rgb' should set args.preview_method to Latent2RGB."""
        set_preview_method("latent2rgb")
        assert args.preview_method == LatentPreviewMethod.Latent2RGB

    def test_override_with_auto(self):
        """'auto' should set args.preview_method to Auto."""
        set_preview_method("auto")
        assert args.preview_method == LatentPreviewMethod.Auto

    def test_override_with_none_value(self):
        """'none' should set args.preview_method to NoPreviews."""
        set_preview_method("none")
        assert args.preview_method == LatentPreviewMethod.NoPreviews

    def test_default_restores_original(self):
        """'default' should restore to default_preview_method."""
        # First override to something else
        set_preview_method("taesd")
        assert args.preview_method == LatentPreviewMethod.TAESD

        # Then use 'default' to restore
        set_preview_method("default")
        assert args.preview_method == default_preview_method

    def test_none_param_restores_original(self):
        """None parameter should restore to default_preview_method."""
        # First override to something else
        set_preview_method("taesd")
        assert args.preview_method == LatentPreviewMethod.TAESD

        # Then use None to restore
        set_preview_method(None)
        assert args.preview_method == default_preview_method

    def test_empty_string_restores_original(self):
        """Empty string should restore to default_preview_method."""
        set_preview_method("taesd")
        set_preview_method("")
        assert args.preview_method == default_preview_method

    def test_invalid_value_restores_original(self):
        """Invalid value should restore to default_preview_method."""
        set_preview_method("taesd")
        set_preview_method("invalid_method")
        assert args.preview_method == default_preview_method

    def test_case_sensitive_invalid_restores(self):
        """Case-mismatched values should restore to default."""
        set_preview_method("taesd")
        set_preview_method("TAESD")  # Wrong case
        assert args.preview_method == default_preview_method


class TestDefaultPreviewMethod:
    """Test default_preview_method module variable."""

    def test_default_is_not_none(self):
        """default_preview_method should not be None."""
        assert default_preview_method is not None

    def test_default_is_enum_member(self):
        """default_preview_method should be a LatentPreviewMethod enum."""
        assert isinstance(default_preview_method, LatentPreviewMethod)

    def test_default_matches_args_initial(self):
        """default_preview_method should match CLI default or user setting."""
        # This tests that default_preview_method was captured at module load
        # After set_preview_method(None), args should equal default
        original = args.preview_method
        set_preview_method("taesd")
        set_preview_method(None)
        assert args.preview_method == default_preview_method
        args.preview_method = original


class TestArgsPreviewMethodModification:
    """Test args.preview_method can be modified correctly."""

    def setup_method(self):
        """Store original value before each test."""
        self.original = args.preview_method

    def teardown_method(self):
        """Restore original value after each test."""
        args.preview_method = self.original

    def test_args_accepts_all_enum_values(self):
        """args.preview_method should accept all LatentPreviewMethod values."""
        for method in LatentPreviewMethod:
            args.preview_method = method
            assert args.preview_method == method

    def test_args_modification_and_restoration(self):
        """args.preview_method should be modifiable and restorable."""
        original = args.preview_method

        args.preview_method = LatentPreviewMethod.TAESD
        assert args.preview_method == LatentPreviewMethod.TAESD

        args.preview_method = original
        assert args.preview_method == original


class TestExecutionFlow:
    """Test the execution flow pattern used in execution.py."""

    def setup_method(self):
        """Store original value before each test."""
        self.original = args.preview_method

    def teardown_method(self):
        """Restore original value after each test."""
        args.preview_method = self.original

    def test_sequential_executions_with_different_methods(self):
        """Simulate multiple queue executions with different preview methods."""
        # Execution 1: taesd
        set_preview_method("taesd")
        assert args.preview_method == LatentPreviewMethod.TAESD

        # Execution 2: none
        set_preview_method("none")
        assert args.preview_method == LatentPreviewMethod.NoPreviews

        # Execution 3: default (restore)
        set_preview_method("default")
        assert args.preview_method == default_preview_method

        # Execution 4: auto
        set_preview_method("auto")
        assert args.preview_method == LatentPreviewMethod.Auto

        # Execution 5: no override (None)
        set_preview_method(None)
        assert args.preview_method == default_preview_method

    def test_override_then_default_pattern(self):
        """Test the pattern: override -> execute -> next call restores."""
        # First execution with override
        set_preview_method("latent2rgb")
        assert args.preview_method == LatentPreviewMethod.Latent2RGB

        # Second execution without override restores default
        set_preview_method(None)
        assert args.preview_method == default_preview_method

    def test_extra_data_simulation(self):
        """Simulate extra_data.get('preview_method') patterns."""
        # Simulate: extra_data = {"preview_method": "taesd"}
        extra_data = {"preview_method": "taesd"}
        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == LatentPreviewMethod.TAESD

        # Simulate: extra_data = {}
        extra_data = {}
        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == default_preview_method

        # Simulate: extra_data = {"preview_method": "default"}
        extra_data = {"preview_method": "default"}
        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == default_preview_method


class TestRealWorldScenarios:
    """Tests using real-world prompt data patterns."""

    def setup_method(self):
        """Store original value before each test."""
        self.original = args.preview_method

    def teardown_method(self):
        """Restore original value after each test."""
        args.preview_method = self.original

    def test_captured_prompt_without_preview_method(self):
        """
        Test with captured prompt that has no preview_method.
        Based on: tests-unit/execution_test/fixtures/default_prompt.json
        """
        # Real captured extra_data structure (preview_method absent)
        extra_data = {
            "extra_pnginfo": {"workflow": {}},
            "client_id": "271314f0dabd48e5aaa488ed7a4ceb0d",
            "create_time": 1765416558179
        }

        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == default_preview_method

    def test_captured_prompt_with_preview_method_taesd(self):
        """Test captured prompt with preview_method: taesd."""
        extra_data = {
            "extra_pnginfo": {"workflow": {}},
            "client_id": "271314f0dabd48e5aaa488ed7a4ceb0d",
            "preview_method": "taesd"
        }

        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == LatentPreviewMethod.TAESD

    def test_captured_prompt_with_preview_method_none(self):
        """Test captured prompt with preview_method: none (disable preview)."""
        extra_data = {
            "extra_pnginfo": {"workflow": {}},
            "client_id": "test-client",
            "preview_method": "none"
        }

        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == LatentPreviewMethod.NoPreviews

    def test_captured_prompt_with_preview_method_latent2rgb(self):
        """Test captured prompt with preview_method: latent2rgb."""
        extra_data = {
            "extra_pnginfo": {"workflow": {}},
            "client_id": "test-client",
            "preview_method": "latent2rgb"
        }

        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == LatentPreviewMethod.Latent2RGB

    def test_captured_prompt_with_preview_method_auto(self):
        """Test captured prompt with preview_method: auto."""
        extra_data = {
            "extra_pnginfo": {"workflow": {}},
            "client_id": "test-client",
            "preview_method": "auto"
        }

        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == LatentPreviewMethod.Auto

    def test_captured_prompt_with_preview_method_default(self):
        """Test captured prompt with preview_method: default (use CLI setting)."""
        # First set to something else
        set_preview_method("taesd")
        assert args.preview_method == LatentPreviewMethod.TAESD

        # Then simulate a prompt with "default"
        extra_data = {
            "extra_pnginfo": {"workflow": {}},
            "client_id": "test-client",
            "preview_method": "default"
        }

        set_preview_method(extra_data.get("preview_method"))
        assert args.preview_method == default_preview_method

    def test_sequential_queue_with_different_preview_methods(self):
        """
        Simulate real queue scenario: multiple prompts with different settings.
        This tests the actual usage pattern in ComfyUI.
        """
        # Queue 1: User wants TAESD preview
        extra_data_1 = {"client_id": "client-1", "preview_method": "taesd"}
        set_preview_method(extra_data_1.get("preview_method"))
        assert args.preview_method == LatentPreviewMethod.TAESD

        # Queue 2: User wants no preview (faster execution)
        extra_data_2 = {"client_id": "client-2", "preview_method": "none"}
        set_preview_method(extra_data_2.get("preview_method"))
        assert args.preview_method == LatentPreviewMethod.NoPreviews

        # Queue 3: User doesn't specify (use server default)
        extra_data_3 = {"client_id": "client-3"}
        set_preview_method(extra_data_3.get("preview_method"))
        assert args.preview_method == default_preview_method

        # Queue 4: User explicitly wants default
        extra_data_4 = {"client_id": "client-4", "preview_method": "default"}
        set_preview_method(extra_data_4.get("preview_method"))
        assert args.preview_method == default_preview_method

        # Queue 5: User wants latent2rgb
        extra_data_5 = {"client_id": "client-5", "preview_method": "latent2rgb"}
        set_preview_method(extra_data_5.get("preview_method"))
        assert args.preview_method == LatentPreviewMethod.Latent2RGB