113 lines
3.2 KiB
Python
113 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
|
|
|
|
"""Prove collision-safe Comfy MODEL callback registration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
import torch
|
|
from comfy.patcher_extension import CallbacksMP
|
|
|
|
from simple_syrup.runtime.model_patcher_mutations import ModelKeyedCallbackMutation
|
|
|
|
|
|
def test_keyed_callback_mutation_installs_exact_callback_on_real_patcher() -> None:
|
|
"""Register through Comfy's public keyed surface without changing identity."""
|
|
|
|
model = _patcher()
|
|
|
|
def callback(patcher: object, unpatch_all: bool) -> None:
|
|
"""Accept the installed detach contract."""
|
|
|
|
del patcher, unpatch_all
|
|
|
|
ModelKeyedCallbackMutation(
|
|
CallbacksMP.ON_DETACH,
|
|
"simple_syrup.device_cache",
|
|
callback,
|
|
).apply(model)
|
|
|
|
assert model.get_callbacks(
|
|
CallbacksMP.ON_DETACH,
|
|
"simple_syrup.device_cache",
|
|
) == [callback]
|
|
|
|
|
|
def test_keyed_callback_mutation_rejects_collision_before_mutating() -> None:
|
|
"""Keep one authoritative owner for each namespaced callback key."""
|
|
|
|
model = _patcher()
|
|
|
|
def first(*args: object) -> tuple[object, ...]:
|
|
"""Return the first callback arguments."""
|
|
|
|
return args
|
|
|
|
def second(*args: object) -> tuple[object, ...]:
|
|
"""Return the second callback arguments."""
|
|
|
|
return args
|
|
|
|
mutation = ModelKeyedCallbackMutation(
|
|
CallbacksMP.ON_DETACH,
|
|
"simple_syrup.device_cache",
|
|
first,
|
|
)
|
|
mutation.apply(model)
|
|
|
|
with pytest.raises(ValueError, match="already installed"):
|
|
ModelKeyedCallbackMutation(
|
|
CallbacksMP.ON_DETACH,
|
|
"simple_syrup.device_cache",
|
|
second,
|
|
).apply(model)
|
|
|
|
assert model.get_callbacks(
|
|
CallbacksMP.ON_DETACH,
|
|
"simple_syrup.device_cache",
|
|
) == [first]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("callback_type", "key", "callback", "error_type", "message"),
|
|
[
|
|
("", "simple_syrup.valid", lambda: None, ValueError, "non-empty"),
|
|
("event", "foreign.valid", lambda: None, ValueError, "must start"),
|
|
("event", "simple_syrup.", lambda: None, ValueError, "namespaced"),
|
|
("event", "simple_syrup.valid", None, TypeError, "must be callable"),
|
|
],
|
|
)
|
|
def test_keyed_callback_mutation_validates_owned_fields(
|
|
callback_type: str,
|
|
key: str,
|
|
callback: Callable[..., object] | None,
|
|
error_type: type[Exception],
|
|
message: str,
|
|
) -> None:
|
|
"""Reject malformed declarations before inspecting Comfy callback state."""
|
|
|
|
model = _patcher()
|
|
|
|
with pytest.raises(error_type, match=message):
|
|
ModelKeyedCallbackMutation(
|
|
callback_type,
|
|
key,
|
|
cast(Callable[..., object], callback),
|
|
).apply(model)
|
|
|
|
assert model.get_all_callbacks(callback_type) == []
|
|
|
|
|
|
def _patcher() -> Any:
|
|
"""Create a real CPU Comfy MODEL patcher."""
|
|
|
|
from comfy.model_patcher import ModelPatcher
|
|
|
|
device = torch.device("cpu")
|
|
return ModelPatcher(torch.nn.Linear(2, 2), device, device)
|