from pathlib import Path
import tempfile
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from diffusion_cli.config import (
DEFAULT_HEIGHT,
DEFAULT_OUTPUT_EXTENSION,
DEFAULT_OUTPUT_QUALITY,
DEFAULT_WIDTH,
GenerationDefaults,
ModelPathConfig,
UserConfig,
TOKENIZER_FILES,
buildDefaultGenerationRequest,
buildGenerationConfig,
loadUserConfig,
normalizeOutputExtension,
outputMimeType,
validateDimensions,
validateOutputQuality,
validateOutputPath,
validateTokenizerPath,
)
from diffusion_cli.errors import DiffusionCliError
class ConfigTest(unittest.TestCase):
def testDefaultDimensionsAreValid(self):
validateDimensions(DEFAULT_WIDTH, DEFAULT_HEIGHT)
def testInvalidWidthFails(self):
with self.assertRaises(DiffusionCliError) as context:
validateDimensions(830, DEFAULT_HEIGHT)
self.assertIn("Width must be", str(context.exception))
def testTokenizerValidation(self):
with tempfile.TemporaryDirectory() as temp_dir:
path = Path(temp_dir)
for file_name in TOKENIZER_FILES:
(path / file_name).write_text("{}", encoding="utf-8")
self.assertEqual(validateTokenizerPath(path), path.resolve())
def testOutputParentIsCreated(self):
with tempfile.TemporaryDirectory() as temp_dir:
output = Path(temp_dir) / "nested" / "image.png"
result = validateOutputPath(output)
self.assertEqual(result, output)
def testOutputExtensionValidationNormalizesAliases(self):
self.assertEqual(normalizeOutputExtension("jpg"), "jpg")
self.assertEqual(normalizeOutputExtension(".JPEG"), "jpg")
self.assertEqual(normalizeOutputExtension(" webp"), "webp")
self.assertEqual(outputMimeType("jpg"), "image/jpeg")
def testOutputExtensionValidationRejectsUnsupportedValues(self):
with self.assertRaises(DiffusionCliError) as context:
normalizeOutputExtension("gif")
self.assertIn(
"Output extension must be one of png, jpg, webp, avif: got gif",
str(context.exception),
)
def testOutputQualityValidationRejectsOutOfRangeAndBool(self):
self.assertEqual(validateOutputQuality(1), 1)
self.assertEqual(validateOutputQuality(100), 100)
with self.assertRaises(DiffusionCliError) as context:
validateOutputQuality(0)
self.assertIn(
"Output quality must be between 1 and 100",
str(context.exception),
)
with self.assertRaises(DiffusionCliError) as context:
validateOutputQuality(True)
self.assertIn(
"Output quality must be an integer",
str(context.exception),
)
def testMissingUserConfigReturnsEmptyConfig(self):
with tempfile.TemporaryDirectory() as temp_dir:
config = loadUserConfig(Path(temp_dir) / "missing.toml")
self.assertIsNone(config.models.diffusion_model)
self.assertIsNone(config.generation.steps)
def testLoadUserConfig(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
config_path = root / "diffusion.toml"
config_path.write_text(
"\n".join(
(
"[models]",
f'checkpoint = "{root / "aio.safetensors"}"',
f'diffusion_model = "{root / "diffusion.safetensors"}"',
f'text_encoder = "{root / "text_encoder.safetensors"}"',
f'vae = "{root / "ae.safetensors"}"',
f'tokenizer = "{root / "tokenizer"}"',
"",
"[generation]",
'negative_prompt = "low quality"',
"width = 512",
"height = 768",
"batch_size = 2",
"steps = 12",
"cfg = 1.5",
f'output = "{root / "output.png"}"',
'output_extension = "jpeg"',
"output_quality = 85",
'device = "cuda:0"',
'dtype = "bf16"',
)
),
encoding="utf-8",
)
config = loadUserConfig(config_path)
self.assertEqual(
config.models.checkpoint.name,
"aio.safetensors",
)
self.assertEqual(
config.models.diffusion_model.name,
"diffusion.safetensors",
)
self.assertEqual(config.models.tokenizer.name, "tokenizer")
self.assertEqual(config.generation.width, 512)
self.assertEqual(config.generation.height, 768)
self.assertEqual(config.generation.cfg, 1.5)
self.assertEqual(config.generation.output_extension, "jpeg")
self.assertEqual(config.generation.output_quality, 85)
self.assertEqual(config.generation.dtype, "bf16")
def testMalformedUserConfigFailsClearly(self):
with tempfile.TemporaryDirectory() as temp_dir:
config_path = Path(temp_dir) / "diffusion.toml"
config_path.write_text("[generation\n", encoding="utf-8")
with self.assertRaises(DiffusionCliError) as context:
loadUserConfig(config_path)
self.assertIn("Invalid TOML config", str(context.exception))
self.assertIn("diffusion.toml", str(context.exception))
def testUnknownTopLevelConfigTableFails(self):
with tempfile.TemporaryDirectory() as temp_dir:
config_path = Path(temp_dir) / "diffusion.toml"
config_path.write_text("[model]\n", encoding="utf-8")
with self.assertRaises(DiffusionCliError) as context:
loadUserConfig(config_path)
self.assertIn(
"Unknown config key top-level.model",
str(context.exception),
)
def testUnknownModelConfigKeyFails(self):
with tempfile.TemporaryDirectory() as temp_dir:
config_path = Path(temp_dir) / "diffusion.toml"
config_path.write_text(
"[models]\ndiffusion = '/tmp/model.safetensors'\n",
encoding="utf-8",
)
with self.assertRaises(DiffusionCliError) as context:
loadUserConfig(config_path)
self.assertIn(
"Unknown config key models.diffusion",
str(context.exception),
)
def testUnknownGenerationConfigKeyFails(self):
with tempfile.TemporaryDirectory() as temp_dir:
config_path = Path(temp_dir) / "diffusion.toml"
config_path.write_text(
"[generation]\nstep = 10\n",
encoding="utf-8",
)
with self.assertRaises(DiffusionCliError) as context:
loadUserConfig(config_path)
self.assertIn(
"Unknown config key generation.step",
str(context.exception),
)
def testWrongConfigScalarTypeFails(self):
with tempfile.TemporaryDirectory() as temp_dir:
config_path = Path(temp_dir) / "diffusion.toml"
config_path.write_text(
'[generation]\nsteps = "10"\n',
encoding="utf-8",
)
with self.assertRaises(DiffusionCliError) as context:
loadUserConfig(config_path)
self.assertIn(
"generation.steps must be an integer",
str(context.exception),
)
def testWrongCheckpointTypeFails(self):
with tempfile.TemporaryDirectory() as temp_dir:
config_path = Path(temp_dir) / "diffusion.toml"
config_path.write_text(
"[models]\ncheckpoint = 12\n",
encoding="utf-8",
)
with self.assertRaises(DiffusionCliError) as context:
loadUserConfig(config_path)
self.assertIn(
"models.checkpoint must be a string path",
str(context.exception),
)
def testGenerationConfigUsesConfigDefaults(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
tokenizer = root / "tokenizer"
tokenizer.mkdir()
for file_name in TOKENIZER_FILES:
(tokenizer / file_name).write_text("{}", encoding="utf-8")
user_config = UserConfig(
models=ModelPathConfig(tokenizer=tokenizer),
generation=GenerationDefaults(
negative_prompt="low quality",
width=512,
height=768,
batch_size=2,
steps=12,
cfg=1.5,
output=root / "image.png",
output_extension="jpeg",
output_quality=85,
device="cuda:0",
dtype="bf16",
),
)
args = SimpleNamespace(
prompt="a mug",
negative_prompt=None,
seed=123,
width=None,
height=None,
batch_size=None,
steps=None,
cfg=None,
output=None,
device=None,
dtype=None,
tokenizer_path=None,
)
with patch("diffusion_cli.config.selectDevice") as select_device:
with patch("diffusion_cli.config.selectDtype") as select_dtype:
select_device.return_value = "cuda:0"
select_dtype.return_value = "bf16"
config = buildGenerationConfig(args, user_config)
self.assertEqual(config.negative_prompt, "low quality")
self.assertEqual(config.width, 512)
self.assertEqual(config.height, 768)
self.assertEqual(config.batch_size, 2)
self.assertEqual(config.steps, 12)
self.assertEqual(config.cfg, 1.5)
self.assertEqual(config.seed, 123)
self.assertEqual(config.output_extension, "jpg")
self.assertEqual(config.output_quality, 85)
self.assertEqual(config.output_mime_type, "image/jpeg")
def testGenerationConfigCliOverridesConfig(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
tokenizer = root / "tokenizer"
tokenizer.mkdir()
for file_name in TOKENIZER_FILES:
(tokenizer / file_name).write_text("{}", encoding="utf-8")
user_config = UserConfig(
models=ModelPathConfig(tokenizer=tokenizer),
generation=GenerationDefaults(
width=512,
steps=12,
output_extension="jpg",
output_quality=75,
),
)
args = SimpleNamespace(
prompt="a mug",
negative_prompt=None,
seed=123,
width=640,
height=None,
batch_size=None,
steps=4,
cfg=None,
output=root / "image.png",
output_extension="webp",
output_quality=80,
device="cuda",
dtype="fp32",
tokenizer_path=None,
)
with patch("diffusion_cli.config.selectDevice") as select_device:
with patch("diffusion_cli.config.selectDtype") as select_dtype:
select_device.return_value = "cuda"
select_dtype.return_value = "fp32"
config = buildGenerationConfig(args, user_config)
self.assertEqual(config.width, 640)
self.assertEqual(config.height, DEFAULT_HEIGHT)
self.assertEqual(config.steps, 4)
self.assertEqual(config.output_extension, "webp")
self.assertEqual(config.output_quality, 80)
def testGenerationConfigUsesBuiltInOutputDefaults(self):
with tempfile.TemporaryDirectory() as temp_dir:
root = Path(temp_dir)
tokenizer = root / "tokenizer"
tokenizer.mkdir()
for file_name in TOKENIZER_FILES:
(tokenizer / file_name).write_text("{}", encoding="utf-8")
user_config = UserConfig(
models=ModelPathConfig(tokenizer=tokenizer),
generation=GenerationDefaults(),
)
args = SimpleNamespace(
prompt="a mug",
negative_prompt=None,
seed=123,
width=None,
height=None,
batch_size=None,
steps=None,
cfg=None,
output=root / "image.png",
device="cuda",
dtype="fp32",
tokenizer_path=None,
)
with patch("diffusion_cli.config.selectDevice") as select_device:
with patch("diffusion_cli.config.selectDtype") as select_dtype:
select_device.return_value = "cuda"
select_dtype.return_value = "fp32"
config = buildGenerationConfig(args, user_config)
self.assertEqual(config.output_extension, DEFAULT_OUTPUT_EXTENSION)
self.assertEqual(config.output_quality, DEFAULT_OUTPUT_QUALITY)
def testDefaultGenerationRequestUsesUiVisibleFallbacks(self):
user_config = UserConfig(
models=ModelPathConfig(),
generation=GenerationDefaults(width=512, steps=12),
)
request = buildDefaultGenerationRequest(user_config)
self.assertEqual(request.prompt, "")
self.assertEqual(request.negative_prompt, "")
self.assertEqual(request.width, 512)
self.assertEqual(request.height, DEFAULT_HEIGHT)
self.assertEqual(request.steps, 12)
self.assertEqual(request.seed, -1)
self.assertEqual(request.output_extension, "png")
self.assertEqual(request.output_quality, 95)
if __name__ == "__main__":
unittest.main()