Vito Sansevero
e84ec6721c
- Update pyproject.toml to use black's default line-length of 88 - This matches what the CI workflow expects (black --check without args) - Reformat all Python files to comply with the new line length - This will prevent CI failures due to formatting discrepancies
353 lines
14 KiB
Python
353 lines
14 KiB
Python
"""Tests for Sampler Combo node."""
|
|
|
|
import pytest
|
|
from unittest.mock import patch
|
|
from kikotools.tools.sampler_combo.node import SamplerComboNode
|
|
from kikotools.tools.sampler_combo.logic import (
|
|
validate_sampler_settings,
|
|
get_sampler_combo,
|
|
get_compatible_scheduler_suggestions,
|
|
get_recommended_steps_range,
|
|
get_recommended_cfg_range,
|
|
get_sampler_info,
|
|
SAMPLERS,
|
|
SCHEDULERS,
|
|
)
|
|
|
|
|
|
class TestSamplerComboLogic:
|
|
"""Test cases for sampler combo logic functions."""
|
|
|
|
def test_validate_sampler_settings_valid(self):
|
|
"""Test validation with valid settings."""
|
|
assert validate_sampler_settings("euler", "normal", 20, 7.0) is True
|
|
assert validate_sampler_settings("dpmpp_2m", "karras", 15, 8.5) is True
|
|
assert validate_sampler_settings("ddim", "ddim_uniform", 30, 6.0) is True
|
|
|
|
def test_validate_sampler_settings_invalid_sampler(self):
|
|
"""Test validation with invalid sampler."""
|
|
assert validate_sampler_settings("invalid_sampler", "normal", 20, 7.0) is False
|
|
|
|
def test_validate_sampler_settings_invalid_scheduler(self):
|
|
"""Test validation with invalid scheduler."""
|
|
assert validate_sampler_settings("euler", "invalid_scheduler", 20, 7.0) is False
|
|
|
|
def test_validate_sampler_settings_invalid_steps(self):
|
|
"""Test validation with invalid steps."""
|
|
assert validate_sampler_settings("euler", "normal", 0, 7.0) is False
|
|
assert validate_sampler_settings("euler", "normal", 1001, 7.0) is False
|
|
assert validate_sampler_settings("euler", "normal", -5, 7.0) is False
|
|
|
|
def test_validate_sampler_settings_invalid_cfg(self):
|
|
"""Test validation with invalid CFG."""
|
|
assert validate_sampler_settings("euler", "normal", 20, -1.0) is False
|
|
assert validate_sampler_settings("euler", "normal", 20, 31.0) is False
|
|
|
|
def test_get_sampler_combo_valid(self):
|
|
"""Test getting sampler combo with valid inputs."""
|
|
result = get_sampler_combo("euler", "normal", 20, 7.0)
|
|
assert result == ("euler", "normal", 20, 7.0)
|
|
|
|
result = get_sampler_combo("dpmpp_2m", "karras", 25, 8.5)
|
|
assert result == ("dpmpp_2m", "karras", 25, 8.5)
|
|
|
|
def test_get_sampler_combo_invalid_returns_defaults(self):
|
|
"""Test that invalid inputs return safe defaults."""
|
|
result = get_sampler_combo("invalid", "normal", 20, 7.0)
|
|
assert result == ("euler", "normal", 20, 7.0)
|
|
|
|
result = get_sampler_combo("euler", "invalid", 20, 7.0)
|
|
assert result == ("euler", "normal", 20, 7.0)
|
|
|
|
def test_get_sampler_combo_sanitizes_values(self):
|
|
"""Test that values are sanitized to valid ranges."""
|
|
# Test steps clamping
|
|
result = get_sampler_combo("euler", "normal", 0, 7.0)
|
|
assert result[2] >= 1 # steps should be at least 1
|
|
|
|
result = get_sampler_combo("euler", "normal", 1500, 7.0)
|
|
assert result[2] <= 1000 # steps should be at most 1000
|
|
|
|
# Test CFG clamping
|
|
result = get_sampler_combo("euler", "normal", 20, -5.0)
|
|
assert result[3] >= 0.0 # CFG should be at least 0
|
|
|
|
result = get_sampler_combo("euler", "normal", 20, 50.0)
|
|
assert result[3] <= 30.0 # CFG should be at most 30
|
|
|
|
def test_get_compatible_scheduler_suggestions(self):
|
|
"""Test getting scheduler suggestions for different samplers."""
|
|
suggestions = get_compatible_scheduler_suggestions("euler")
|
|
assert isinstance(suggestions, list)
|
|
assert len(suggestions) > 0
|
|
assert "normal" in suggestions
|
|
|
|
suggestions = get_compatible_scheduler_suggestions("ddim")
|
|
assert "ddim_uniform" in suggestions
|
|
|
|
# Test unknown sampler returns defaults
|
|
suggestions = get_compatible_scheduler_suggestions("unknown_sampler")
|
|
assert "normal" in suggestions
|
|
assert "karras" in suggestions
|
|
|
|
def test_get_recommended_steps_range(self):
|
|
"""Test getting recommended steps range for samplers."""
|
|
min_steps, max_steps, default_steps = get_recommended_steps_range("euler")
|
|
assert isinstance(min_steps, int)
|
|
assert isinstance(max_steps, int)
|
|
assert isinstance(default_steps, int)
|
|
assert min_steps <= default_steps <= max_steps
|
|
assert min_steps > 0
|
|
|
|
# Test unknown sampler returns defaults
|
|
min_steps, max_steps, default_steps = get_recommended_steps_range("unknown")
|
|
assert min_steps == 10
|
|
assert max_steps == 50
|
|
assert default_steps == 20
|
|
|
|
def test_get_recommended_cfg_range(self):
|
|
"""Test getting recommended CFG range for samplers."""
|
|
min_cfg, max_cfg, default_cfg = get_recommended_cfg_range("euler")
|
|
assert isinstance(min_cfg, float)
|
|
assert isinstance(max_cfg, float)
|
|
assert isinstance(default_cfg, float)
|
|
assert min_cfg <= default_cfg <= max_cfg
|
|
assert min_cfg >= 0.0
|
|
|
|
# Test unknown sampler returns defaults
|
|
min_cfg, max_cfg, default_cfg = get_recommended_cfg_range("unknown")
|
|
assert min_cfg == 1.0
|
|
assert max_cfg == 20.0
|
|
assert default_cfg == 7.0
|
|
|
|
def test_get_sampler_info(self):
|
|
"""Test getting sampler information."""
|
|
info = get_sampler_info()
|
|
assert isinstance(info, dict)
|
|
assert "samplers" in info
|
|
assert "schedulers" in info
|
|
assert "sampler_count" in info
|
|
assert "scheduler_count" in info
|
|
assert info["sampler_count"] == len(SAMPLERS)
|
|
assert info["scheduler_count"] == len(SCHEDULERS)
|
|
|
|
|
|
class TestSamplerComboNode:
|
|
"""Test cases for SamplerComboNode."""
|
|
|
|
def setup_method(self):
|
|
"""Set up test fixtures."""
|
|
self.node = SamplerComboNode()
|
|
|
|
def test_input_types_structure(self):
|
|
"""Test that INPUT_TYPES returns correct structure."""
|
|
input_types = SamplerComboNode.INPUT_TYPES()
|
|
|
|
assert "required" in input_types
|
|
required = input_types["required"]
|
|
|
|
# Check all required inputs are present
|
|
assert "sampler_name" in required
|
|
assert "scheduler" in required
|
|
assert "steps" in required
|
|
assert "cfg" in required
|
|
|
|
# Check sampler input structure
|
|
sampler_input = required["sampler_name"]
|
|
assert sampler_input[0] == SAMPLERS
|
|
assert isinstance(sampler_input[1], dict)
|
|
assert "default" in sampler_input[1]
|
|
assert "tooltip" in sampler_input[1]
|
|
|
|
# Check scheduler input structure
|
|
scheduler_input = required["scheduler"]
|
|
assert scheduler_input[0] == SCHEDULERS
|
|
assert isinstance(scheduler_input[1], dict)
|
|
|
|
# Check steps input structure
|
|
steps_input = required["steps"]
|
|
assert steps_input[0] == "INT"
|
|
assert steps_input[1]["min"] == 1
|
|
assert steps_input[1]["max"] == 100
|
|
|
|
# Check CFG input structure
|
|
cfg_input = required["cfg"]
|
|
assert cfg_input[0] == "FLOAT"
|
|
assert cfg_input[1]["min"] == 0.0
|
|
assert cfg_input[1]["max"] == 20.0
|
|
|
|
def test_return_types_structure(self):
|
|
"""Test that return types are correctly defined."""
|
|
assert SamplerComboNode.RETURN_TYPES == ("SAMPLER", SCHEDULERS, "INT", "FLOAT")
|
|
assert SamplerComboNode.RETURN_NAMES == (
|
|
"sampler_name",
|
|
"scheduler",
|
|
"steps",
|
|
"cfg",
|
|
)
|
|
assert SamplerComboNode.FUNCTION == "get_sampler_combo"
|
|
assert SamplerComboNode.CATEGORY == "ComfyAssets"
|
|
|
|
def test_get_sampler_combo_valid_inputs(self):
|
|
"""Test get_sampler_combo with valid inputs."""
|
|
result = self.node.get_sampler_combo("euler", "normal", 20, 7.0)
|
|
assert result == ("euler", "normal", 20, 7.0)
|
|
|
|
result = self.node.get_sampler_combo("dpmpp_2m", "karras", 15, 8.5)
|
|
assert result == ("dpmpp_2m", "karras", 15, 8.5)
|
|
|
|
def test_get_sampler_combo_invalid_inputs_returns_defaults(self):
|
|
"""Test that invalid inputs return safe defaults."""
|
|
with patch.object(self.node, "handle_error") as mock_error:
|
|
mock_error.side_effect = ValueError("Invalid settings")
|
|
|
|
try:
|
|
result = self.node.get_sampler_combo("invalid", "normal", 20, 7.0)
|
|
except ValueError:
|
|
pass # Expected when handle_error raises
|
|
|
|
# Test with exception handling bypassed
|
|
with patch(
|
|
"kikotools.tools.sampler_combo.node.validate_sampler_settings",
|
|
return_value=False,
|
|
):
|
|
result = self.node.get_sampler_combo("invalid", "normal", 20, 7.0)
|
|
assert result == ("euler", "normal", 20, 7.0)
|
|
|
|
def test_validate_inputs_valid(self):
|
|
"""Test input validation with valid inputs."""
|
|
# Should not raise any exception
|
|
self.node.validate_inputs("euler", "normal", 20, 7.0)
|
|
|
|
def test_validate_inputs_invalid(self):
|
|
"""Test input validation with invalid inputs."""
|
|
with pytest.raises(ValueError):
|
|
self.node.validate_inputs("invalid", "normal", 20, 7.0)
|
|
|
|
def test_get_scheduler_suggestions(self):
|
|
"""Test getting scheduler suggestions."""
|
|
suggestions = self.node.get_scheduler_suggestions("euler")
|
|
assert isinstance(suggestions, list)
|
|
assert len(suggestions) > 0
|
|
|
|
suggestions = self.node.get_scheduler_suggestions("ddim")
|
|
assert "ddim_uniform" in suggestions
|
|
|
|
def test_get_steps_recommendation(self):
|
|
"""Test getting steps recommendations."""
|
|
rec = self.node.get_steps_recommendation("euler")
|
|
assert isinstance(rec, dict)
|
|
assert "min" in rec
|
|
assert "max" in rec
|
|
assert "default" in rec
|
|
assert "recommendation" in rec
|
|
|
|
def test_get_cfg_recommendation(self):
|
|
"""Test getting CFG recommendations."""
|
|
rec = self.node.get_cfg_recommendation("euler")
|
|
assert isinstance(rec, dict)
|
|
assert "min" in rec
|
|
assert "max" in rec
|
|
assert "default" in rec
|
|
assert "recommendation" in rec
|
|
|
|
def test_get_combo_analysis(self):
|
|
"""Test getting combo analysis."""
|
|
analysis = self.node.get_combo_analysis("euler", "normal", 20, 7.0)
|
|
assert isinstance(analysis, dict)
|
|
assert "sampler" in analysis
|
|
assert "scheduler" in analysis
|
|
assert "steps" in analysis
|
|
assert "cfg" in analysis
|
|
assert "valid" in analysis
|
|
assert "scheduler_suggestions" in analysis
|
|
assert "scheduler_compatible" in analysis
|
|
assert "steps_optimal" in analysis
|
|
assert "cfg_optimal" in analysis
|
|
|
|
def test_get_available_samplers(self):
|
|
"""Test getting available samplers."""
|
|
samplers = SamplerComboNode.get_available_samplers()
|
|
assert isinstance(samplers, list)
|
|
assert len(samplers) > 0
|
|
assert "euler" in samplers
|
|
|
|
def test_get_available_schedulers(self):
|
|
"""Test getting available schedulers."""
|
|
schedulers = SamplerComboNode.get_available_schedulers()
|
|
assert isinstance(schedulers, list)
|
|
assert len(schedulers) > 0
|
|
assert "normal" in schedulers
|
|
|
|
def test_string_representations(self):
|
|
"""Test string representations of the node."""
|
|
str_repr = str(self.node)
|
|
assert "SamplerComboNode" in str_repr
|
|
assert "samplers=" in str_repr
|
|
assert "schedulers=" in str_repr
|
|
|
|
repr_str = repr(self.node)
|
|
assert "SamplerComboNode" in repr_str
|
|
assert "category=" in repr_str
|
|
assert "function=" in repr_str
|
|
|
|
def test_node_inheritance(self):
|
|
"""Test that node properly inherits from base class."""
|
|
from kikotools.base.base_node import ComfyAssetsBaseNode
|
|
|
|
assert isinstance(self.node, ComfyAssetsBaseNode)
|
|
assert hasattr(self.node, "validate_inputs")
|
|
assert hasattr(self.node, "handle_error")
|
|
assert hasattr(self.node, "log_info")
|
|
|
|
|
|
class TestSamplerComboIntegration:
|
|
"""Integration tests for Sampler Combo functionality."""
|
|
|
|
def test_full_workflow_valid_settings(self):
|
|
"""Test complete workflow with valid settings."""
|
|
node = SamplerComboNode()
|
|
|
|
# Test with different sampler/scheduler combinations
|
|
test_cases = [
|
|
("euler", "normal", 20, 7.0),
|
|
("dpmpp_2m", "karras", 15, 8.0),
|
|
("euler_ancestral", "exponential", 25, 9.0),
|
|
("ddim", "ddim_uniform", 30, 6.0),
|
|
]
|
|
|
|
for sampler, scheduler, steps, cfg in test_cases:
|
|
result = node.get_sampler_combo(sampler, scheduler, steps, cfg)
|
|
assert result == (sampler, scheduler, steps, cfg)
|
|
|
|
def test_recommendation_compatibility(self):
|
|
"""Test that recommendations are compatible with actual functionality."""
|
|
node = SamplerComboNode()
|
|
|
|
for sampler in SAMPLERS[:5]: # Test first 5 samplers
|
|
suggestions = node.get_scheduler_suggestions(sampler)
|
|
steps_rec = node.get_steps_recommendation(sampler)
|
|
cfg_rec = node.get_cfg_recommendation(sampler)
|
|
|
|
# Test that recommendations work with the node
|
|
for scheduler in suggestions[:2]: # Test first 2 suggestions
|
|
result = node.get_sampler_combo(
|
|
sampler, scheduler, steps_rec["default"], cfg_rec["default"]
|
|
)
|
|
assert result[0] == sampler
|
|
assert result[1] == scheduler
|
|
assert result[2] == steps_rec["default"]
|
|
assert result[3] == cfg_rec["default"]
|
|
|
|
def test_error_recovery(self):
|
|
"""Test error recovery with malformed inputs."""
|
|
node = SamplerComboNode()
|
|
|
|
# These should all return safe defaults due to error handling
|
|
with patch(
|
|
"kikotools.tools.sampler_combo.logic.validate_sampler_settings",
|
|
side_effect=Exception("Simulated error"),
|
|
):
|
|
result = node.get_sampler_combo("euler", "normal", 20, 7.0)
|
|
assert result == ("euler", "normal", 20, 7.0) # Safe defaults
|