103 lines
3.2 KiB
Python
103 lines
3.2 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Tests for the Simple Load Checkpoint Comfy v3 wrapper."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from types import ModuleType
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from simple_syrup.nodes.simple_load_checkpoint import SimpleLoadCheckpoint
|
|
from simple_syrup.nodes_v3.simple_load_checkpoint import SimpleLoadCheckpointV3
|
|
from simple_syrup.runtime.checkpoint_loader import (
|
|
CLIP_SKIP_DEFAULT,
|
|
USE_CHECKPOINT_VAE_CHOICE,
|
|
)
|
|
|
|
|
|
class FakeFolderPaths(ModuleType):
|
|
"""Small folder_paths fake for v3 schema declarations."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Create deterministic ComfyUI filename lists."""
|
|
|
|
super().__init__("folder_paths")
|
|
self.files = {
|
|
"checkpoints": ["model.safetensors"],
|
|
"vae": ["manual_vae.safetensors"],
|
|
"vae_approx": [],
|
|
}
|
|
|
|
def get_filename_list(self, folder_name: str) -> list[str]:
|
|
"""Return deterministic filenames for a model folder."""
|
|
|
|
return self.files[folder_name]
|
|
|
|
|
|
def test_simple_load_checkpoint_v3_schema(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""The v3 schema mirrors the legacy checkpoint loader contract."""
|
|
|
|
monkeypatch.setitem(sys.modules, "folder_paths", FakeFolderPaths())
|
|
|
|
schema = SimpleLoadCheckpointV3.define_schema()
|
|
|
|
assert schema.node_id == "SimpleSyrup.SimpleLoadCheckpoint"
|
|
assert schema.display_name == "Simple Load Checkpoint"
|
|
assert [input_item.id for input_item in schema.inputs] == [
|
|
"ckpt_name",
|
|
"vae_name",
|
|
"clip_skip",
|
|
]
|
|
assert schema.inputs[0].options == ["model.safetensors"]
|
|
assert schema.inputs[1].options == [
|
|
USE_CHECKPOINT_VAE_CHOICE,
|
|
"manual_vae.safetensors",
|
|
"pixel_space",
|
|
]
|
|
assert schema.inputs[1].default == USE_CHECKPOINT_VAE_CHOICE
|
|
assert schema.inputs[2].io_type == "BOOLEAN"
|
|
assert schema.inputs[2].default is CLIP_SKIP_DEFAULT
|
|
assert [output.id for output in schema.outputs] == ["model", "clip", "vae"]
|
|
assert [output.io_type for output in schema.outputs] == ["MODEL", "CLIP", "VAE"]
|
|
|
|
|
|
def test_simple_load_checkpoint_v3_execute_forwards_to_legacy_loader(
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""The v3 wrapper forwards execution to the legacy loader."""
|
|
|
|
class FakeService:
|
|
"""Service double for the legacy loader."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize captured kwargs."""
|
|
|
|
self.kwargs: dict[str, object] = {}
|
|
|
|
def load_checkpoint(self, **kwargs: object) -> tuple[str, str, str]:
|
|
"""Return fixed checkpoint outputs."""
|
|
|
|
self.kwargs = kwargs
|
|
return ("model", "clip", "vae")
|
|
|
|
fake_service = FakeService()
|
|
monkeypatch.setattr(SimpleLoadCheckpoint, "_service", fake_service)
|
|
|
|
assert SimpleLoadCheckpointV3.execute(
|
|
"model.safetensors",
|
|
USE_CHECKPOINT_VAE_CHOICE,
|
|
True,
|
|
) == ("model", "clip", "vae")
|
|
assert fake_service.kwargs == {
|
|
"ckpt_name": "model.safetensors",
|
|
"vae_name": USE_CHECKPOINT_VAE_CHOICE,
|
|
"clip_skip": True,
|
|
}
|