"""Tests for workflow template builders."""
from unittest.mock import patch
from src.workflows import (
build_img2img_workflow,
build_txt2img_workflow,
get_default_negative_prompt,
)
# --- get_default_negative_prompt ---
def test_default_negative_prompt_is_non_empty() -> None:
result = get_default_negative_prompt()
assert isinstance(result, str)
assert len(result) > 0
# --- build_txt2img_workflow structure ---
REQUIRED_TXT2IMG_NODE_IDS = {"3", "4", "5", "6", "7", "8", "9"}
REQUIRED_TXT2IMG_CLASS_TYPES = {
"3": "KSampler",
"4": "CheckpointLoaderSimple",
"5": "EmptyLatentImage",
"6": "CLIPTextEncode",
"7": "CLIPTextEncode",
"8": "VAEDecode",
"9": "SaveImage",
}
@patch("src.workflows.settings")
def test_txt2img_workflow_structure(mock_settings: object) -> None:
mock_settings.default_model = "test.safetensors"
mock_settings.default_width = 512
mock_settings.default_height = 512
mock_settings.default_steps = 10
mock_settings.default_cfg_scale = 5.0
workflow = build_txt2img_workflow(prompt="a cat")
# All required nodes exist
assert set(workflow.keys()) == REQUIRED_TXT2IMG_NODE_IDS
# Each node has correct class_type
for node_id, expected_class in REQUIRED_TXT2IMG_CLASS_TYPES.items():
assert workflow[node_id]["class_type"] == expected_class
# Each node has an "inputs" dict
for node_id in REQUIRED_TXT2IMG_NODE_IDS:
assert "inputs" in workflow[node_id]
@patch("src.workflows.settings")
def test_txt2img_workflow_connections(mock_settings: object) -> None:
mock_settings.default_model = "test.safetensors"
mock_settings.default_width = 512
mock_settings.default_height = 512
mock_settings.default_steps = 10
mock_settings.default_cfg_scale = 5.0
workflow = build_txt2img_workflow(prompt="test")
ksampler = workflow["3"]["inputs"]
# Model comes from checkpoint loader (node 4, output 0)
assert ksampler["model"] == ["4", 0]
# Positive conditioning from CLIP encode (node 6, output 0)
assert ksampler["positive"] == ["6", 0]
# Negative conditioning from CLIP encode (node 7, output 0)
assert ksampler["negative"] == ["7", 0]
# Latent image from EmptyLatentImage (node 5, output 0)
assert ksampler["latent_image"] == ["5", 0]
# CLIP text encoders reference checkpoint's CLIP output (node 4, output 1)
assert workflow["6"]["inputs"]["clip"] == ["4", 1]
assert workflow["7"]["inputs"]["clip"] == ["4", 1]
# VAEDecode takes samples from KSampler (node 3, output 0)
assert workflow["8"]["inputs"]["samples"] == ["3", 0]
# VAEDecode takes VAE from checkpoint (node 4, output 2)
assert workflow["8"]["inputs"]["vae"] == ["4", 2]
# SaveImage takes images from VAEDecode (node 8, output 0)
assert workflow["9"]["inputs"]["images"] == ["8", 0]
@patch("src.workflows.settings")
def test_txt2img_default_params(mock_settings: object) -> None:
mock_settings.default_model = "default_model.safetensors"
mock_settings.default_width = 1024
mock_settings.default_height = 1024
mock_settings.default_steps = 20
mock_settings.default_cfg_scale = 7.0
workflow = build_txt2img_workflow(prompt="a dog")
# Model defaults to settings
assert workflow["4"]["inputs"]["ckpt_name"] == "default_model.safetensors"
# Dimensions default to settings
assert workflow["5"]["inputs"]["width"] == 1024
assert workflow["5"]["inputs"]["height"] == 1024
# Steps and cfg default to settings
assert workflow["3"]["inputs"]["steps"] == 20
assert workflow["3"]["inputs"]["cfg"] == 7.0
# Seed defaults to -1
assert workflow["3"]["inputs"]["seed"] == -1
# Prompt text is set correctly
assert workflow["6"]["inputs"]["text"] == "a dog"
# Negative prompt falls back to default negative prompt
assert workflow["7"]["inputs"]["text"] == get_default_negative_prompt()
@patch("src.workflows.settings")
def test_txt2img_custom_params(mock_settings: object) -> None:
mock_settings.default_model = "default.safetensors"
mock_settings.default_width = 1024
mock_settings.default_height = 1024
mock_settings.default_steps = 20
mock_settings.default_cfg_scale = 7.0
workflow = build_txt2img_workflow(
prompt="a landscape",
negative_prompt="ugly, blurry",
model="custom_model.safetensors",
width=768,
height=768,
steps=30,
cfg_scale=9.5,
seed=42,
)
assert workflow["4"]["inputs"]["ckpt_name"] == "custom_model.safetensors"
assert workflow["5"]["inputs"]["width"] == 768
assert workflow["5"]["inputs"]["height"] == 768
assert workflow["3"]["inputs"]["steps"] == 30
assert workflow["3"]["inputs"]["cfg"] == 9.5
assert workflow["3"]["inputs"]["seed"] == 42
assert workflow["6"]["inputs"]["text"] == "a landscape"
assert workflow["7"]["inputs"]["text"] == "ugly, blurry"
# --- build_img2img_workflow ---
REQUIRED_IMG2IMG_NODE_IDS = {"3", "4", "5", "6", "7", "8", "9", "10"}
REQUIRED_IMG2IMG_CLASS_TYPES = {
"3": "KSampler",
"4": "CheckpointLoaderSimple",
"5": "LoadImage",
"6": "CLIPTextEncode",
"7": "CLIPTextEncode",
"8": "VAEDecode",
"9": "SaveImage",
"10": "VAEEncode",
}
@patch("src.workflows.settings")
def test_img2img_workflow_structure(mock_settings: object) -> None:
mock_settings.default_model = "test.safetensors"
mock_settings.default_steps = 10
mock_settings.default_cfg_scale = 5.0
workflow = build_img2img_workflow(prompt="enhance", image_path="input.png")
assert set(workflow.keys()) == REQUIRED_IMG2IMG_NODE_IDS
for node_id, expected_class in REQUIRED_IMG2IMG_CLASS_TYPES.items():
assert workflow[node_id]["class_type"] == expected_class
@patch("src.workflows.settings")
def test_img2img_uses_load_image_node(mock_settings: object) -> None:
mock_settings.default_model = "test.safetensors"
mock_settings.default_steps = 10
mock_settings.default_cfg_scale = 5.0
workflow = build_img2img_workflow(prompt="test", image_path="my_image.png")
assert workflow["5"]["class_type"] == "LoadImage"
assert workflow["5"]["inputs"]["image"] == "my_image.png"
@patch("src.workflows.settings")
def test_img2img_ksampler_uses_vae_encode(mock_settings: object) -> None:
mock_settings.default_model = "test.safetensors"
mock_settings.default_steps = 10
mock_settings.default_cfg_scale = 5.0
workflow = build_img2img_workflow(prompt="test", image_path="input.png")
# KSampler latent_image should come from VAEEncode (node 10), not EmptyLatentImage
assert workflow["3"]["inputs"]["latent_image"] == ["10", 0]
# VAEEncode takes pixels from LoadImage (node 5)
assert workflow["10"]["inputs"]["pixels"] == ["5", 0]
@patch("src.workflows.settings")
def test_img2img_denoise_parameter(mock_settings: object) -> None:
mock_settings.default_model = "test.safetensors"
mock_settings.default_steps = 10
mock_settings.default_cfg_scale = 5.0
# Default denoise
workflow = build_img2img_workflow(prompt="test", image_path="input.png")
assert workflow["3"]["inputs"]["denoise"] == 0.75
# Custom denoise
workflow = build_img2img_workflow(prompt="test", image_path="input.png", denoise=0.5)
assert workflow["3"]["inputs"]["denoise"] == 0.5
