from __future__ import annotations import threading import time import os import sys import tempfile import types import unittest from pathlib import Path from unittest import mock import yaml from core.runtime_config import CONFIG, estimate_gpu_duration from core.settings import CHECKPOINT_DIR, INPUT_DIR, OUTPUT_DIR from core.task_scheduler import ( QueueFullError, TaskCancelledError, generation_guard, generation_slot, submit_background, ) ROOT = Path(__file__).resolve().parents[1] class RuntimeConfigTests(unittest.TestCase): def test_duration_estimation_is_bounded(self): self.assertEqual(estimate_gpu_duration({"zero_gpu_duration": 999}), 120) self.assertEqual(estimate_gpu_duration({"zero_gpu_duration": 1}), 30) self.assertEqual( estimate_gpu_duration( { "model_display_name": "example-lightning", "num_inference_steps": 4, "batch_size": 1, "width": 1024, "height": 1024, } ), 45, ) self.assertEqual( estimate_gpu_duration( { "model_display_name": "large-model", "num_inference_steps": 40, "batch_size": 3, "width": 2048, "height": 2048, } ), 120, ) def test_generation_slots_respect_configured_limit(self): active = 0 maximum = 0 state_lock = threading.Lock() def worker(): nonlocal active, maximum with generation_slot(): with state_lock: active += 1 maximum = max(maximum, active) time.sleep(0.02) with state_lock: active -= 1 threads = [threading.Thread(target=worker) for _ in range(CONFIG.gpu_concurrency + 2)] for thread in threads: thread.start() for thread in threads: thread.join() self.assertLessEqual(maximum, CONFIG.gpu_concurrency) def test_cancelled_job_stops_before_guarded_execution(self): reached_function = False @generation_guard def guarded(ui_inputs): nonlocal reached_function reached_function = True cancel_event = threading.Event() cancel_event.set() with self.assertRaises(TaskCancelledError): guarded({"_cancel_event": cancel_event}) self.assertFalse(reached_function) def test_waiting_cancelled_job_leaves_gate_before_gpu_is_free(self): holder_started = threading.Event() release_holder = threading.Event() cancellation = threading.Event() errors = [] def holder(): with generation_slot(): holder_started.set() release_holder.wait(2) @generation_guard def waiting_job(ui_inputs): raise AssertionError("cancelled job must not execute") def wait_then_cancel(): try: waiting_job({"_cancel_event": cancellation}) except BaseException as exc: errors.append(exc) holder_thread = threading.Thread(target=holder) holder_thread.start() self.assertTrue(holder_started.wait(1)) waiting_thread = threading.Thread(target=wait_then_cancel) waiting_thread.start() cancellation.set() waiting_thread.join(1) self.assertFalse(waiting_thread.is_alive()) self.assertEqual(len(errors), 1) self.assertIsInstance(errors[0], TaskCancelledError) release_holder.set() holder_thread.join(1) self.assertFalse(holder_thread.is_alive()) def test_mcp_pending_queue_is_bounded(self): release = threading.Event() futures = [ submit_background(lambda: release.wait(2)) for _ in range(CONFIG.mcp_max_pending) ] with self.assertRaises(QueueFullError): submit_background(lambda: None) release.set() for future in futures: self.assertTrue(future.result(timeout=3)) class RegistryTests(unittest.TestCase): def test_runtime_directories_are_project_absolute(self): for configured in (CHECKPOINT_DIR, INPUT_DIR, OUTPUT_DIR): path = Path(configured) self.assertTrue(path.is_absolute()) self.assertTrue(path.is_relative_to(ROOT)) def test_quick_presets_exist(self): registry = yaml.safe_load((ROOT / "yaml" / "model_list.yaml").read_text("utf-8")) names = { model["display_name"] for architecture in registry["Checkpoint"].values() for model in architecture.get("models", []) } expected = { "Krea-2-Turbo", "lightx2v/Qwen-Image-2512-Lightning", "circlestone-labs/Anima-Turbo-v1.0", "lightx2v/Qwen-Image-Edit-2511-Lightning", "CagliostroLab/Animagine XL 4.0", } self.assertTrue(expected.issubset(names)) def test_vendor_revisions_are_full_commits(self): lock = yaml.safe_load((ROOT / "vendor.lock.yaml").read_text("utf-8")) entries = [lock["comfyui"], *lock["custom_nodes"].values()] for entry in entries: revision = entry["revision"] self.assertEqual(len(revision), 40) int(revision, 16) def test_task_input_recipes_are_complete(self): input_dir = ROOT / "core" / "pipelines" / "workflow_recipes" / "_partials" / "input" task_recipes = { "txt2img": "txt2img_latent.yaml", "img2img": "img2img.yaml", "inpaint": "inpaint.yaml", "outpaint": "outpaint.yaml", "hires_fix": "hires_fix.yaml", } for task_type, recipe_name in task_recipes.items(): recipe = yaml.safe_load((input_dir / recipe_name).read_text("utf-8")) self.assertIn( "latent_source", recipe.get("nodes", {}), f"{task_type} must provide the sampler latent_source", ) txt2img_router = yaml.safe_load((input_dir / "txt2img.yaml").read_text("utf-8")) self.assertEqual( txt2img_router["imports"], ["txt2img_{{ latent_type }}.yaml"] ) def test_concurrency_regressions_are_absent(self): mcp_run = (ROOT / "mcp_tools" / "run.py").read_text("utf-8") input_processor = ( ROOT / "core" / "pipelines" / "pipeline_input_processor.py" ).read_text("utf-8") studio = (ROOT / "ui" / "shared" / "studio_ui.py").read_text("utf-8") requirements = (ROOT / "requirements.txt").read_text("utf-8") self.assertNotIn("threading.Thread", mcp_run) self.assertIn("uuid.uuid4().hex", input_processor) self.assertIn('"_task_prefixes": [(prefix, None)]', studio) self.assertIn("onnxruntime-gpu==", requirements) class ComfySetupTests(unittest.TestCase): def test_initialize_registers_application_model_directories(self): from comfy_integration import setup with tempfile.TemporaryDirectory() as temp_dir: root = Path(temp_dir) comfyui_path = root / "ComfyUI" comfyui_path.mkdir() (comfyui_path / "nodes.py").touch() model_dir = root / "models" / "checkpoints" input_dir = root / "input" output_dir = root / "output" folder_paths = types.ModuleType("folder_paths") folder_paths.add_model_folder_path = mock.Mock() folder_paths.set_input_directory = mock.Mock() folder_paths.set_output_directory = mock.Mock() comfy_package = types.ModuleType("comfy") model_management = types.ModuleType("comfy.model_management") comfy_package.model_management = model_management with ( mock.patch.object( setup, "CATEGORY_TO_DIR_MAP", {"checkpoints": str(model_dir)} ), mock.patch.object(setup, "INPUT_DIR", str(input_dir)), mock.patch.object(setup, "OUTPUT_DIR", str(output_dir)), mock.patch.object(setup, "_load_lock", return_value={"comfyui": {}}), mock.patch.dict( os.environ, { "COMFYUI_PATH": str(comfyui_path), "IMAGEGEN_SKIP_CUSTOM_NODES": "1", }, clear=False, ), mock.patch.dict( sys.modules, { "folder_paths": folder_paths, "comfy": comfy_package, "comfy.model_management": model_management, }, ), ): setup.initialize_comfyui() folder_paths.add_model_folder_path.assert_called_once_with( "checkpoints", str(model_dir.resolve()), is_default=True ) folder_paths.set_input_directory.assert_called_once_with( str(input_dir.resolve()) ) folder_paths.set_output_directory.assert_called_once_with( str(output_dir.resolve()) ) if __name__ == "__main__": unittest.main()