zlei6/OpenMontage
1
0
Fork 0
mirror of https://github.com/calesthio/OpenMontage.git synced 2026-08-24 17:10:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into codex/repair-pr-457

Browse Source 
# Conflicts:
#	tools/graphics/image_selector.py
This commit is contained in:
calesthio
2026-08-13 10:17:11 -07:00
parent e614c9d4c4 6e957f2378
commit 4e5bd8d5fa
16 changed files with 2740 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

641
tests/contracts/test_comfyui_tools.py
View File

@@ -17,13 +17,14 @@ from tools.base_tool import (
ToolStatus,
ToolTier,
)
from tools.audio.comfyui_music import ComfyUIMusic
from tools.graphics.comfyui_image import ComfyUIImage
from tools.graphics.image_selector import ImageSelector
from tools.tool_registry import ToolRegistry
from tools.video.video_selector import VideoSelector
from tools.video.comfyui_video import ComfyUIVideo
TOOLS = [ComfyUIImage, ComfyUIVideo]
TOOLS = [ComfyUIImage, ComfyUIVideo, ComfyUIMusic]
WORKFLOW_DIR = Path(__file__).resolve().parent.parent.parent / "tools" / "_comfyui" / "workflows"
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
@@ -136,6 +137,7 @@ EXPECTED_WORKFLOWS = [
"flux2-txt2img.json",
"wan22-i2v-4step.json",
"wan22-t2v-4step.json",
"ace-step-1-t2a.json",
]
@@ -282,6 +284,73 @@ class TestClientHelpers:
"folder_type": "temp",
}
def test_poll_timeout_carries_prompt_id_for_recovery(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient, ComfyUIError
client = ComfyUIClient("http://comfy.test")
monkeypatch.setattr(
"tools._comfyui.client.requests.get",
lambda *a, **k: type("R", (), {
"raise_for_status": lambda self: None,
"json": lambda self: {},
})(),
)
monkeypatch.setattr("tools._comfyui.client.time.sleep", lambda s: None)
with pytest.raises(ComfyUIError) as excinfo:
client.poll("prompt-timeout-1", timeout=0, interval=0)
assert excinfo.value.prompt_id == "prompt-timeout-1"
assert "prompt-timeout-1" in str(excinfo.value)
def test_generate_resume_prompt_id_skips_resubmit(self, monkeypatch, tmp_path):
from tools._comfyui.client import ComfyUIClient
import sys
client = ComfyUIClient("http://comfy.test")
def fail_submit(workflow):
raise AssertionError("submit() should not be called when resuming")
monkeypatch.setattr(client, "submit", fail_submit)
_install_fake_websocket(monkeypatch, frames=[])
monkeypatch.setattr(
sys.modules["websocket"],
"create_connection",
lambda *a, **k: (_ for _ in ()).throw(
AssertionError("resumed jobs must use history polling")
),
)
monkeypatch.setattr(client, "poll", lambda prompt_id, **kwargs: {
"outputs": {"9": {"images": [{
"filename": "resumed.png", "subfolder": "", "type": "output",
}]}}
})
monkeypatch.setattr(client, "download", lambda filename, subfolder, dest, folder_type="output": Path(dest))
paths = client.generate(
{"9": {"inputs": {}}}, "9", tmp_path / "out.png",
resume_prompt_id="already-running-id",
)
assert paths == [tmp_path / "out.png"]
def test_generate_reads_audio_key_from_savaudio_node(self, monkeypatch, tmp_path):
"""The native SaveAudio node writes outputs under "audio", not
"images"/"gifs" -- comfyui_music depends on this being handled."""
from tools._comfyui.client import ComfyUIClient
client = ComfyUIClient("http://comfy.test")
monkeypatch.setattr(client, "submit", lambda workflow: "p1")
monkeypatch.setattr(client, "poll", lambda prompt_id, **kwargs: {
"outputs": {"9": {"audio": [{
"filename": "track.flac", "subfolder": "", "type": "output",
}]}}
})
monkeypatch.setattr(client, "download", lambda filename, subfolder, dest, folder_type="output": Path(dest))
paths = client.generate({"9": {"inputs": {}}}, "9", tmp_path / "out.flac")
assert paths == [tmp_path / "out.flac"]
def test_is_default_url_when_env_not_set(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
@@ -310,6 +379,290 @@ class TestClientHelpers:
assert "myhost:9999" in msg
assert "COMFYUI_SERVER_URL" not in msg
def test_submit_includes_client_id_for_websocket_targeting(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
client = ComfyUIClient("http://comfy.test")
seen = {}
def fake_post(url, json=None, timeout=None):
seen.update(json)
return type("R", (), {
"raise_for_status": lambda self: None,
"json": lambda self: {"prompt_id": "abc"},
})()
monkeypatch.setattr("tools._comfyui.client.requests.post", fake_post)
client.submit({"1": {"inputs": {}}})
assert seen["client_id"] == client.client_id
class TestMultiServer:
def test_capability_env_var_takes_priority_over_shared(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
client = ComfyUIClient(capability="video")
assert client.server_url == "http://video-gpu:8188"
def test_falls_back_to_shared_when_capability_var_unset(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
monkeypatch.delenv("COMFYUI_IMAGE_SERVER_URL", raising=False)
client = ComfyUIClient(capability="image")
assert client.server_url == "http://shared:8188"
def test_other_capability_env_var_does_not_leak_across_tools(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
monkeypatch.delenv("COMFYUI_VIDEO_SERVER_URL", raising=False)
video_client = ComfyUIClient(capability="video")
assert video_client.server_url == "http://localhost:8188"
def test_explicit_server_url_wins_over_capability_env_var(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
client = ComfyUIClient("http://explicit:1234", capability="video")
assert client.server_url == "http://explicit:1234"
def test_no_capability_behaves_as_before(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
client = ComfyUIClient()
assert client.server_url == "http://shared:8188"
assert client.is_default_url is False
def test_is_default_url_true_only_when_both_vars_unset(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
monkeypatch.delenv("COMFYUI_IMAGE_SERVER_URL", raising=False)
client = ComfyUIClient(capability="image")
assert client.is_default_url is True
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
client2 = ComfyUIClient(capability="image")
assert client2.is_default_url is False
def test_unavailable_reason_mentions_capability_and_shared_var(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
monkeypatch.delenv("COMFYUI_VIDEO_SERVER_URL", raising=False)
client = ComfyUIClient(capability="video")
msg = client.unavailable_reason()
assert "COMFYUI_VIDEO_SERVER_URL" in msg
assert "COMFYUI_SERVER_URL" in msg
def test_image_and_video_tools_use_independent_servers(self, monkeypatch):
from tools.graphics.comfyui_image import ComfyUIImage
from tools.video.comfyui_video import ComfyUIVideo
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
image_tool = ComfyUIImage()
video_tool = ComfyUIVideo()
assert image_tool._client.server_url == "http://image-gpu:8188"
assert video_tool._client.server_url == "http://video-gpu:8188"
class _FakeWSTimeout(Exception):
pass
class _FakeWSConn:
def __init__(self, frames):
self._frames = list(frames)
def settimeout(self, value):
pass
def recv(self):
if not self._frames:
raise _FakeWSTimeout()
return self._frames.pop(0)
def close(self):
pass
def _install_fake_websocket(monkeypatch, frames):
"""Inject a fake `websocket` module so wait_ws() runs without the real
optional websocket-client dependency installed."""
import sys
import types
fake_module = types.SimpleNamespace(
WebSocketTimeoutException=_FakeWSTimeout,
create_connection=lambda url, timeout=10: _FakeWSConn(frames),
)
monkeypatch.setitem(sys.modules, "websocket", fake_module)
class TestWebsocketWait:
def test_wait_ws_returns_job_completed_before_connection(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
import sys
client = ComfyUIClient("http://comfy.test")
_install_fake_websocket(monkeypatch, frames=[])
monkeypatch.setattr(
sys.modules["websocket"],
"create_connection",
lambda *a, **k: (_ for _ in ()).throw(
AssertionError("completed history must avoid websocket connection")
),
)
monkeypatch.setattr(
"tools._comfyui.client.requests.get",
lambda *a, **k: type("R", (), {
"raise_for_status": lambda self: None,
"json": lambda self: {"done": {"outputs": {"9": {}}}},
})(),
)
assert client.wait_ws("done", timeout=5) == {"outputs": {"9": {}}}
def test_wait_ws_completes_on_executing_none_node(self, monkeypatch, tmp_path):
from tools._comfyui.client import ComfyUIClient
client = ComfyUIClient("http://comfy.test")
progress_events = []
frames = [
json.dumps({"type": "progress", "data": {
"value": 2, "max": 20, "prompt_id": "p1",
}}),
json.dumps({"type": "executing", "data": {
"node": None, "prompt_id": "p1",
}}),
]
_install_fake_websocket(monkeypatch, frames)
history_calls = iter(({}, {}, {"p1": {"outputs": {"9": {}}}}))
monkeypatch.setattr(
"tools._comfyui.client.requests.get",
lambda *a, **k: type("R", (), {
"raise_for_status": lambda self: None,
"json": lambda self: next(history_calls),
})(),
)
entry = client.wait_ws("p1", timeout=5, on_progress=progress_events.append)
assert entry == {"outputs": {"9": {}}}
assert progress_events == [{"value": 2, "max": 20, "prompt_id": "p1"}]
def test_wait_ws_execution_error_raises_with_prompt_id(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient, ComfyUIError
client = ComfyUIClient("http://comfy.test")
frames = [
json.dumps({"type": "execution_error", "data": {
"prompt_id": "p2", "exception_message": "boom",
}}),
]
_install_fake_websocket(monkeypatch, frames)
with pytest.raises(ComfyUIError) as excinfo:
client.wait_ws("p2", timeout=5)
assert excinfo.value.prompt_id == "p2"
def test_wait_ws_ignores_other_prompts_on_shared_connection(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient
client = ComfyUIClient("http://comfy.test")
frames = [
# Another job's event on the same client_id -- must not trigger completion.
json.dumps({"type": "executing", "data": {
"node": None, "prompt_id": "someone-elses-job",
}}),
json.dumps({"type": "executing", "data": {
"node": None, "prompt_id": "p3",
}}),
]
_install_fake_websocket(monkeypatch, frames)
monkeypatch.setattr(
"tools._comfyui.client.requests.get",
lambda *a, **k: type("R", (), {
"raise_for_status": lambda self: None,
"json": lambda self: {"p3": {"outputs": {}}},
})(),
)
entry = client.wait_ws("p3", timeout=5)
assert entry == {"outputs": {}}
def test_wait_ws_timeout_raises_comfyuierror_with_prompt_id(self, monkeypatch):
from tools._comfyui.client import ComfyUIClient, ComfyUIError
client = ComfyUIClient("http://comfy.test")
_install_fake_websocket(monkeypatch, frames=[]) # recv() always times out
with pytest.raises(ComfyUIError) as excinfo:
client.wait_ws("p4", timeout=0)
assert excinfo.value.prompt_id == "p4"
def test_wait_falls_back_to_poll_when_websocket_unavailable(self, monkeypatch):
"""No websocket-client installed (or any transport failure) must
silently fall back to REST polling, not blow up the whole call."""
from tools._comfyui.client import ComfyUIClient
import sys
client = ComfyUIClient("http://comfy.test")
monkeypatch.delitem(sys.modules, "websocket", raising=False)
monkeypatch.setattr(
"builtins.__import__",
_raise_on_websocket_import(__import__),
)
monkeypatch.setattr(
client, "poll", lambda prompt_id, **kwargs: {"outputs": {"used": "poll"}}
)
entry = client._wait("p5", timeout=5, interval=5)
assert entry == {"outputs": {"used": "poll"}}
def test_wait_does_not_swallow_genuine_comfyuierror_from_websocket(self, monkeypatch):
"""A real execution error detected over the websocket must propagate,
not be masked by a fallback-to-poll retry."""
from tools._comfyui.client import ComfyUIClient, ComfyUIError
client = ComfyUIClient("http://comfy.test")
frames = [
json.dumps({"type": "execution_error", "data": {
"prompt_id": "p6", "exception_message": "bad node",
}}),
]
_install_fake_websocket(monkeypatch, frames)
def fail_poll(prompt_id, **kwargs):
raise AssertionError("poll() should not be called after a real ws error")
monkeypatch.setattr(client, "poll", fail_poll)
with pytest.raises(ComfyUIError) as excinfo:
client._wait("p6", timeout=5, interval=5)
assert excinfo.value.prompt_id == "p6"
def _raise_on_websocket_import(real_import):
def _import(name, *args, **kwargs):
if name == "websocket":
raise ImportError("no module named websocket")
return real_import(name, *args, **kwargs)
return _import
# ------------------------------------------------------------------
# Model discovery (offline, no server needed)
@@ -332,6 +685,11 @@ class TestModelRequirements:
assert len(_REQUIRED_MODELS_T2V) > 0
assert any("t2v" in m.lower() for m in _REQUIRED_MODELS_T2V)
def test_music_tool_has_required_models(self):
from tools.audio.comfyui_music import _REQUIRED_MODELS
assert len(_REQUIRED_MODELS) > 0
assert any("ace_step" in m.lower() for m in _REQUIRED_MODELS)
# ------------------------------------------------------------------
# Custom workflow contract and provenance
@@ -418,6 +776,77 @@ class TestCustomWorkflowContract:
assert provenance["model_stack"] == [{"role": "lora", "name": "style.safetensors"}]
assert provenance["model_stack_source"] == "caller_supplied"
def test_video_timeout_surfaces_resumable_prompt_id(self, tmp_path):
from tools._comfyui.client import ComfyUIError
tool = ComfyUIVideo()
tool._client.is_available = lambda: True
def fake_generate(workflow, output_node, dest, **kwargs):
raise ComfyUIError("Prompt timed-out-id did not complete within 5s", prompt_id="timed-out-id")
tool._client.generate = fake_generate
result = tool.execute({
"prompt": "test",
"workflow_json": json.dumps({"42": {"inputs": {}}}),
"output_node": "42",
"output_path": str(tmp_path / "video.mp4"),
"timeout_seconds": 5,
})
assert result.success is False
assert result.data["prompt_id"] == "timed-out-id"
assert "resume_prompt_id" in result.error
assert "timed-out-id" in result.error
def test_video_passes_timeout_and_resume_prompt_id_through(self, tmp_path):
tool = ComfyUIVideo()
tool._client.is_available = lambda: True
seen = {}
def fake_generate(workflow, output_node, dest, **kwargs):
seen.update(kwargs)
return [Path(dest)]
tool._client.generate = fake_generate
result = tool.execute({
"prompt": "test",
"workflow_json": json.dumps({"42": {"inputs": {}}}),
"output_node": "42",
"output_path": str(tmp_path / "video.mp4"),
"timeout_seconds": 7200,
"resume_prompt_id": "already-running-id",
})
assert result.success is True
assert seen["timeout"] == 7200
assert seen["resume_prompt_id"] == "already-running-id"
def test_video_default_timeout_is_generous_not_900s(self, tmp_path):
tool = ComfyUIVideo()
tool._client.is_available = lambda: True
seen = {}
def fake_generate(workflow, output_node, dest, **kwargs):
seen.update(kwargs)
return [Path(dest)]
tool._client.generate = fake_generate
tool.execute({
"prompt": "test",
"workflow_json": json.dumps({"42": {"inputs": {}}}),
"output_node": "42",
"output_path": str(tmp_path / "video.mp4"),
})
# Regression guard: the old hardcoded 900s timeout false-failed real
# renders on modest local GPUs (observed ~1360-1630s for non-accelerated
# custom Wan 1.3B workflows at 832x480/81-97 frames).
assert seen["timeout"] > 900
def test_image_missing_models_are_structured(self):
tool = ComfyUIImage()
tool._client.is_available = lambda: True
@@ -466,6 +895,216 @@ class TestCustomWorkflowContract:
assert any(item["role"] == "vae" for item in provenance["model_stack"])
class TestComfyUIMusic:
def test_capability_and_provider(self):
tool = ComfyUIMusic()
assert tool.capability == "music_generation"
assert tool.provider == "comfyui"
def test_bundled_path_requires_no_workflow_json_or_output_node(self, tmp_path):
"""Without workflow_json/workflow_path it should attempt the bundled
ACE-Step workflow, not demand a custom one."""
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
tool._client.check_models = lambda required: (list(required), [])
tool._client.generate = lambda workflow, output_node, dest, **kwargs: [Path(dest)]
result = tool.execute({
"prompt": "ambient pad",
"output_path": str(tmp_path / "music.mp3"),
})
assert result.success is True
assert result.data["workflow_provenance"]["source"] == "bundled"
def test_custom_workflow_without_output_node_errors(self):
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
result = tool.execute({
"prompt": "ambient pad",
"workflow_json": json.dumps({"9": {"inputs": {}}}),
})
assert result.success is False
assert "output_node" in result.error
def test_bundled_missing_models_returns_structured_payload(self):
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
tool._client.check_models = lambda required: ([], list(required))
result = tool.execute({"prompt": "ambient pad"})
assert result.success is False
assert result.data["missing_models"][0]["name"] == "ace_step_v1_3.5b.safetensors"
assert result.data["missing_models"][0]["download_url"]
def test_bundled_generation_patches_tags_lyrics_and_seed(self, tmp_path):
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
tool._client.check_models = lambda required: (list(required), [])
seen = {}
def fake_generate(workflow, output_node, dest, **kwargs):
seen["workflow"] = workflow
seen["output_node"] = output_node
return [Path(dest)]
tool._client.generate = fake_generate
result = tool.execute({
"prompt": "lofi hip hop, chill, rain sounds",
"lyrics": "[verse]\nquiet streets",
"duration_seconds": 45,
"seed": 777,
"output_path": str(tmp_path / "music.mp3"),
})
assert result.success is True
assert seen["output_node"] == "10"
assert seen["workflow"]["2"]["inputs"]["tags"] == "lofi hip hop, chill, rain sounds"
assert seen["workflow"]["2"]["inputs"]["lyrics"] == "[verse]\nquiet streets"
assert seen["workflow"]["4"]["inputs"]["seconds"] == 45
assert seen["workflow"]["8"]["inputs"]["seed"] == 777
assert result.data["model"] == "ace-step-v1-3.5b"
def test_bundled_generation_preserves_seed_zero(self, tmp_path):
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
tool._client.check_models = lambda required: (list(required), [])
seen = {}
def fake_generate(workflow, output_node, dest, **kwargs):
seen["seed"] = workflow["8"]["inputs"]["seed"]
return [Path(dest)]
tool._client.generate = fake_generate
result = tool.execute({
"prompt": "deterministic test",
"seed": 0,
"output_path": str(tmp_path / "music.mp3"),
})
assert result.success is True
assert result.seed == 0
assert seen["seed"] == 0
def test_get_status_degraded_when_model_missing(self):
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
tool._client.check_models = lambda required: ([], list(required))
assert tool.get_status() == ToolStatus.DEGRADED
def test_unavailable_server_reports_unavailable_reason(self):
tool = ComfyUIMusic()
tool._client.is_available = lambda: False
tool._client.unavailable_reason = lambda: "no server here"
result = tool.execute({
"prompt": "ambient pad",
"workflow_json": json.dumps({"9": {"inputs": {}}}),
"output_node": "9",
})
assert result.success is False
assert result.error == "no server here"
def test_successful_generation_returns_provenance_and_duration(self, tmp_path, monkeypatch):
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
dest_file = tmp_path / "music.mp3"
def fake_generate(workflow, output_node, dest, **kwargs):
Path(dest).write_bytes(b"fake-audio-bytes")
return [Path(dest)]
tool._client.generate = fake_generate
monkeypatch.setattr("shutil.which", lambda name: None) # no ffprobe in test env
result = tool.execute({
"prompt": "upbeat synthwave",
"workflow_json": json.dumps({"9": {"inputs": {}}}),
"output_node": "9",
"output_path": str(dest_file),
"workflow_name": "my-ace-step-graph",
"workflow_model": "ace-step-v1-3.5b",
})
assert result.success is True
assert result.data["provider"] == "comfyui"
assert result.data["model"] == "ace-step-v1-3.5b"
assert result.data["output"] == str(dest_file)
assert result.data["format"] == "mp3"
assert result.data["duration_seconds"] is None # ffprobe unavailable
provenance = result.data["workflow_provenance"]
assert provenance["source"] == "user_supplied"
assert provenance["output_node"] == "9"
assert provenance["workflow_hash_sha256"]
def test_timeout_surfaces_resumable_prompt_id(self, tmp_path):
from tools._comfyui.client import ComfyUIError
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
def fake_generate(workflow, output_node, dest, **kwargs):
raise ComfyUIError("timed out", prompt_id="music-prompt-id")
tool._client.generate = fake_generate
result = tool.execute({
"prompt": "ambient pad",
"workflow_json": json.dumps({"9": {"inputs": {}}}),
"output_node": "9",
"output_path": str(tmp_path / "music.mp3"),
})
assert result.success is False
assert result.data["prompt_id"] == "music-prompt-id"
assert "resume_prompt_id" in result.error
def test_passes_timeout_and_resume_prompt_id_through(self, tmp_path):
tool = ComfyUIMusic()
tool._client.is_available = lambda: True
seen = {}
def fake_generate(workflow, output_node, dest, **kwargs):
seen.update(kwargs)
return [Path(dest)]
tool._client.generate = fake_generate
tool.execute({
"prompt": "ambient pad",
"workflow_json": json.dumps({"9": {"inputs": {}}}),
"output_node": "9",
"output_path": str(tmp_path / "music.mp3"),
"timeout_seconds": 3600,
"resume_prompt_id": "already-running-id",
})
assert seen["timeout"] == 3600
assert seen["resume_prompt_id"] == "already-running-id"
def test_registry_discovers_comfyui_music_under_music_generation(self):
registry = ToolRegistry()
tool = ComfyUIMusic()
registry.register(tool)
registry._discovered_packages.add("tools")
by_capability = registry.get_by_capability("music_generation")
assert any(t.name == "comfyui_music" for t in by_capability)
def test_uses_music_capability_env_var_for_multi_server(self, monkeypatch):
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
monkeypatch.setenv("COMFYUI_MUSIC_SERVER_URL", "http://music-gpu:8188")
tool = ComfyUIMusic()
assert tool._client.server_url == "http://music-gpu:8188"
class TestComfyUISetupOffer:
def test_provider_menu_summary_includes_structured_setup_offer(self):

291
tests/tools/test_minimax_image.py Normal file
View File

@@ -0,0 +1,291 @@
"""Contract tests for the MiniMax image generation tool."""
from __future__ import annotations
import base64
import pytest
import requests
from tools.base_tool import ToolStatus
from tools.graphics import minimax_image
from tools.graphics.minimax_image import MiniMaxImage
from tools.tool_registry import ToolRegistry
class FakeResponse:
def __init__(self, *, json_data=None, content: bytes = b"") -> None:
self._json_data = json_data
self.content = content
def raise_for_status(self) -> None:
return None
def json(self):
return self._json_data
@pytest.fixture(autouse=True)
def clear_minimax_env(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("MINIMAX_API_KEY", raising=False)
monkeypatch.delenv("MINIMAX_REGION", raising=False)
monkeypatch.delenv("MINIMAX_BASE_URL", raising=False)
def test_registry_registers_minimax_image_tool() -> None:
registry = ToolRegistry()
assert registry.register_module(minimax_image) == ["minimax_image"]
tool = registry.get("minimax_image")
assert tool is not None
assert tool.provider == "minimax"
assert tool.capability == "image_generation"
assert tool.input_schema["properties"]["model"]["enum"] == [
"image-01",
"image-01-live",
]
def test_status_requires_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
tool = MiniMaxImage()
assert tool.get_status() == ToolStatus.UNAVAILABLE
monkeypatch.setenv("MINIMAX_API_KEY", "test-key")
assert tool.get_status() == ToolStatus.AVAILABLE
def test_cost_estimate_and_result_report_paid_images(
monkeypatch: pytest.MonkeyPatch, tmp_path
) -> None:
monkeypatch.setenv("MINIMAX_API_KEY", "test-key")
monkeypatch.setattr(
requests,
"post",
lambda *args, **kwargs: FakeResponse(
json_data={
"data": {
"image_base64": [
base64.b64encode(b"one").decode("ascii"),
base64.b64encode(b"two").decode("ascii"),
]
},
"base_resp": {"status_code": 0},
}
),
)
tool = MiniMaxImage()
inputs = {
"prompt": "A lighthouse at dusk",
"response_format": "base64",
"n": 2,
"output_path": str(tmp_path / "image.png"),
}
assert tool.estimate_cost(inputs) == pytest.approx(0.007)
result = tool.execute(inputs)
assert result.success, result.error
assert result.cost_usd == pytest.approx(0.007)
def test_image_selector_can_route_to_minimax(
monkeypatch: pytest.MonkeyPatch, tmp_path
) -> None:
from tools.graphics.image_selector import ImageSelector
monkeypatch.setenv("MINIMAX_API_KEY", "test-key")
monkeypatch.setattr(
requests,
"post",
lambda *args, **kwargs: FakeResponse(
json_data={
"data": {
"image_base64": [base64.b64encode(b"image").decode("ascii")]
},
"base_resp": {"status_code": 0},
}
),
)
tool = MiniMaxImage()
selector = ImageSelector()
monkeypatch.setattr(selector, "_providers", lambda: [tool])
result = selector.execute(
{
"prompt": "A lighthouse at dusk",
"preferred_provider": "minimax",
"response_format": "base64",
"output_path": str(tmp_path / "selected.png"),
}
)
assert result.success, result.error
assert result.data["selected_provider"] == "minimax"
assert result.data["selected_tool"] == "minimax_image"
@pytest.mark.parametrize(
("region", "expected_base_url"),
[
("global", "https://api.minimax.io"),
("global_en", "https://api.minimax.io"),
("cn", "https://api.minimaxi.com"),
("cn_zh", "https://api.minimaxi.com"),
],
)
def test_region_routes_to_official_endpoint(
monkeypatch: pytest.MonkeyPatch, region: str, expected_base_url: str
) -> None:
monkeypatch.setenv("MINIMAX_REGION", region)
assert MiniMaxImage()._base_url() == expected_base_url
def test_url_response_downloads_all_images(
monkeypatch: pytest.MonkeyPatch, tmp_path
) -> None:
monkeypatch.setenv("MINIMAX_API_KEY", "test-key")
monkeypatch.setenv("MINIMAX_REGION", "cn")
captured = {}
def fake_post(url, *, headers, json, timeout):
captured.update(url=url, headers=headers, payload=json, timeout=timeout)
return FakeResponse(
json_data={
"id": "request-1",
"data": {
"image_urls": [
"https://example.test/one.png",
"https://example.test/two.png",
]
},
"metadata": {"success_count": 2, "failed_count": 0},
"base_resp": {"status_code": 0, "status_msg": "success"},
}
)
def fake_get(url, *, timeout):
assert timeout == 120
return FakeResponse(content=url.rsplit("/", 1)[-1].encode())
monkeypatch.setattr(requests, "post", fake_post)
monkeypatch.setattr(requests, "get", fake_get)
output_path = tmp_path / "image.png"
result = MiniMaxImage().execute(
{
"prompt": "A lighthouse at dusk",
"model": "image-01-live",
"subject_reference": [
{"type": "character", "image_file": "https://example.test/ref.png"}
],
"aspect_ratio": "16:9",
"response_format": "url",
"seed": 42,
"n": 2,
"prompt_optimizer": True,
"output_path": str(output_path),
}
)
assert result.success
assert captured["url"] == "https://api.minimaxi.com/v1/image_generation"
assert captured["headers"]["Authorization"] == "Bearer test-key"
assert captured["payload"] == {
"model": "image-01-live",
"prompt": "A lighthouse at dusk",
"response_format": "url",
"n": 2,
"prompt_optimizer": True,
"subject_reference": [
{"type": "character", "image_file": "https://example.test/ref.png"}
],
"aspect_ratio": "16:9",
"seed": 42,
}
assert result.artifacts == [
str(tmp_path / "image_1.png"),
str(tmp_path / "image_2.png"),
]
assert (tmp_path / "image_1.png").read_bytes() == b"one.png"
assert (tmp_path / "image_2.png").read_bytes() == b"two.png"
assert result.data["metadata"] == {"success_count": 2, "failed_count": 0}
def test_base64_response_writes_inline_images(
monkeypatch: pytest.MonkeyPatch, tmp_path
) -> None:
monkeypatch.setenv("MINIMAX_API_KEY", "test-key")
image_bytes = b"inline image"
monkeypatch.setattr(
requests,
"post",
lambda *args, **kwargs: FakeResponse(
json_data={
"data": {
"image_base64": [
"data:image/png;base64,"
+ base64.b64encode(image_bytes).decode("ascii")
]
},
"metadata": {"success_count": "1", "failed_count": "0"},
"base_resp": {"status_code": 0},
}
),
)
monkeypatch.setattr(
requests,
"get",
lambda *args, **kwargs: pytest.fail("base64 output must not be downloaded"),
)
output_path = tmp_path / "inline.png"
result = MiniMaxImage().execute(
{
"prompt": "A paper-cut forest",
"response_format": "base64",
"output_path": str(output_path),
}
)
assert result.success
assert output_path.read_bytes() == image_bytes
assert result.data["response_format"] == "base64"
def test_base_response_error_is_returned(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("MINIMAX_API_KEY", "test-key")
monkeypatch.setattr(
requests,
"post",
lambda *args, **kwargs: FakeResponse(
json_data={
"base_resp": {
"status_code": 1008,
"status_msg": "insufficient balance",
}
}
),
)
result = MiniMaxImage().execute({"prompt": "A mountain cabin"})
assert not result.success
assert result.error == "MiniMax API error 1008: insufficient balance"
def test_width_and_height_must_be_provided_together(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("MINIMAX_API_KEY", "test-key")
monkeypatch.setattr(
requests,
"post",
lambda *args, **kwargs: pytest.fail("invalid inputs must not call the API"),
)
result = MiniMaxImage().execute(
{"prompt": "A mountain cabin", "width": 1024}
)
assert not result.success
assert "width and height must be set together" in (result.error or "")

299
tests/tools/test_seedream_image.py Normal file
View File

@@ -0,0 +1,299 @@
"""Regression tests: seedream_image must return every image it requests and bills for.
Covers:
- Multi-image output: all requested images must be written and returned
- Cost estimation: billed count matches delivered artifacts
- Single-image output: exact output path preserved
- Async polling: COMPLETED / FAILED / CANCELLED / timeout paths
- API key validation: graceful failure when FAL_KEY is unset
"""
import sys
import types
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
class _FakeResponse:
def __init__(self, json_data: dict | None = None, status_code: int = 200, content: bytes = b""):
self._json_data = json_data or {}
self.status_code = status_code
self.content = content
def raise_for_status(self):
if self.status_code >= 400:
import requests
raise requests.HTTPError(response=self)
def json(self):
return self._json_data
def _build_submit_response(request_id: str = "req_123") -> _FakeResponse:
return _FakeResponse({"request_id": request_id})
def _build_status_response(status: str, error: str | None = None) -> _FakeResponse:
data = {"status": status}
if error:
data["error"] = error
return _FakeResponse(data)
def _build_result_response(image_urls: list[str]) -> _FakeResponse:
images = [{"url": url} for url in image_urls]
return _FakeResponse({"images": images})
def _build_image_content(index: int) -> bytes:
return f"SEEDREAM_IMAGE_{index}".encode()
@pytest.fixture
def seedream_tool(monkeypatch):
monkeypatch.setenv("FAL_KEY", "test-fal-key")
from tools.graphics.seedream_image import SeedreamImage
return SeedreamImage()
@pytest.fixture
def mock_requests(monkeypatch):
mock_post = MagicMock()
mock_get = MagicMock()
fake_requests = types.ModuleType("requests")
fake_requests.post = mock_post
fake_requests.get = mock_get
fake_requests.HTTPError = type("HTTPError", (Exception,), {})
monkeypatch.setitem(sys.modules, "requests", fake_requests)
return mock_post, mock_get
def _setup_mock_execution(mock_post, mock_get, num_images: int = 1, status: str = "COMPLETED",
error: str | None = None, extra_gets: list = None):
"""Helper to setup common mock execution flow."""
mock_post.return_value = _build_submit_response()
side_effects = [_build_status_response(status, error)]
if status == "COMPLETED":
urls = [f"http://img.url/{i}" for i in range(num_images)]
side_effects.append(_build_result_response(urls))
side_effects.extend([_FakeResponse(content=_build_image_content(i)) for i in range(num_images)])
elif extra_gets:
side_effects.extend(extra_gets)
mock_get.side_effect = side_effects
# ========== Core Regression Tests ==========
class TestMultiOutputRegression:
def test_all_requested_images_are_written(self, seedream_tool, tmp_path, mock_requests):
mock_post, mock_get = mock_requests
_setup_mock_execution(mock_post, mock_get, num_images=3)
result = seedream_tool.execute({
"prompt": "test", "num_images": 3,
"output_format": "jpeg", "output_path": str(tmp_path / "gen.jpeg"),
})
assert result.success
assert result.data["image_count"] == 3
assert len(result.artifacts) == 3
files = sorted(tmp_path.glob("*.jpeg"))
assert len(files) == 3
contents = {f.read_bytes() for f in files}
assert contents == {b"SEEDREAM_IMAGE_0", b"SEEDREAM_IMAGE_1", b"SEEDREAM_IMAGE_2"}
def test_artifacts_match_billed_count(self, seedream_tool, tmp_path, mock_requests):
mock_post, mock_get = mock_requests
_setup_mock_execution(mock_post, mock_get, num_images=4)
inputs = {"prompt": "t", "num_images": 4, "output_path": str(tmp_path / "out.png")}
result = seedream_tool.execute(inputs)
billed = seedream_tool.estimate_cost(inputs)
assert len(result.artifacts) == 4
assert billed == pytest.approx(0.135 * 4)
class TestSingleOutput:
def test_single_image_keeps_exact_path(self, seedream_tool, tmp_path, mock_requests):
mock_post, mock_get = mock_requests
_setup_mock_execution(mock_post, mock_get, num_images=1)
out = tmp_path / "single.png"
result = seedream_tool.execute({"prompt": "s", "num_images": 1, "output_path": str(out)})
assert result.success
assert result.artifacts == [str(out)]
assert out.read_bytes() == b"SEEDREAM_IMAGE_0"
# ========== Cost Estimation (Parameterized) ==========
class TestCostEstimation:
@pytest.mark.parametrize("size,expected", [
("square", 0.0675), ("landscape_4_3", 0.0675),
("portrait_4_3", 0.0675), ("auto_1K", 0.0675),
])
def test_small_size_pricing(self, seedream_tool, size, expected):
cost = seedream_tool.estimate_cost({"image_size": size, "num_images": 1})
assert cost == pytest.approx(expected)
@pytest.mark.parametrize("size,expected", [
("square_hd", 0.135), ("landscape_16_9", 0.135),
("portrait_16_9", 0.135), ("auto_2K", 0.135),
])
def test_large_size_pricing(self, seedream_tool, size, expected):
cost = seedream_tool.estimate_cost({"image_size": size, "num_images": 1})
assert cost == pytest.approx(expected)
@pytest.mark.parametrize("n", [1, 2, 3, 4])
def test_cost_scales_with_num_images(self, seedream_tool, n):
cost = seedream_tool.estimate_cost({"image_size": "auto_2K", "num_images": n})
assert cost == pytest.approx(round(0.135 * n, 4))
def test_unknown_size_falls_back_to_high_price(self, seedream_tool):
cost = seedream_tool.estimate_cost({"image_size": "unknown", "num_images": 1})
assert cost == pytest.approx(0.135)
def test_default_values(self, seedream_tool):
cost = seedream_tool.estimate_cost({})
assert cost == pytest.approx(0.135)
# ========== Async Polling States ==========
class TestAsyncPolling:
def test_completed_on_first_poll(self, seedream_tool, tmp_path, mock_requests):
mock_post, mock_get = mock_requests
_setup_mock_execution(mock_post, mock_get, num_images=1)
result = seedream_tool.execute({"prompt": "q", "output_path": str(tmp_path / "q.png")})
assert result.success
assert result.data["request_id"]
@pytest.mark.parametrize("status,error_msg", [
("FAILED", "Content policy violation"),
("CANCELLED", None),
])
def test_failed_states_return_error(self, seedream_tool, mock_requests, status, error_msg):
mock_post, mock_get = mock_requests
_setup_mock_execution(mock_post, mock_get, status=status, error=error_msg)
result = seedream_tool.execute({"prompt": "bad"})
assert not result.success
assert status in result.error
def test_timeout_returns_error(self, seedream_tool, mock_requests):
mock_post, mock_get = mock_requests
mock_post.return_value = _build_submit_response()
mock_get.side_effect = [_build_status_response("IN_PROGRESS")] * 100
with patch("tools.graphics.seedream_image.time.sleep"):
result = seedream_tool.execute({"prompt": "timeout"})
assert not result.success
assert "timed out" in result.error.lower()
# ========== Validation & Error Handling ==========
class TestValidation:
@pytest.mark.parametrize("value", [0, 5, 1.5, True])
def test_num_images_rejects_invalid_values(
self, seedream_tool, mock_requests, value
):
mock_post, _ = mock_requests
result = seedream_tool.execute({"prompt": "t", "num_images": value})
assert not result.success
assert "num_images" in (result.error or "")
mock_post.assert_not_called()
def test_missing_api_key_returns_error(self, monkeypatch):
monkeypatch.delenv("FAL_KEY", raising=False)
monkeypatch.delenv("FAL_AI_API_KEY", raising=False)
from tools.graphics.seedream_image import SeedreamImage
result = SeedreamImage().execute({"prompt": "t"})
assert not result.success
assert "FAL_KEY" in result.error
def test_status_available_with_key(self, seedream_tool):
assert seedream_tool.get_status().name == "AVAILABLE"
def test_status_unavailable_without_key(self, monkeypatch):
monkeypatch.delenv("FAL_KEY", raising=False)
monkeypatch.delenv("FAL_AI_API_KEY", raising=False)
from tools.graphics.seedream_image import SeedreamImage
assert SeedreamImage().get_status().name == "UNAVAILABLE"
def test_missing_request_id_raises_error(self, seedream_tool, mock_requests):
mock_post, mock_get = mock_requests
mock_post.return_value = _FakeResponse({})
result = seedream_tool.execute({"prompt": "no id"})
assert not result.success
assert "request_id" in result.error.lower()
def test_completed_without_images_raises_error(self, seedream_tool, mock_requests):
mock_post, mock_get = mock_requests
mock_post.return_value = _build_submit_response()
mock_get.side_effect = [
_build_status_response("COMPLETED"),
_FakeResponse({"images": []}),
]
result = seedream_tool.execute({"prompt": "empty"})
assert not result.success
assert "no images" in result.error.lower()
# ========== Metadata & Integration ==========
class TestMetadata:
def test_provider_and_model_info(self, seedream_tool, tmp_path, mock_requests):
mock_post, mock_get = mock_requests
_setup_mock_execution(mock_post, mock_get, num_images=1)
result = seedream_tool.execute({"prompt": "m", "output_path": str(tmp_path / "m.png")})
assert result.data["provider"] == "seedream"
assert result.data["model"] == "seedream_v5"
assert result.model == "fal-ai/bytedance/seedream/v5"
def test_cost_matches_estimate(self, seedream_tool, tmp_path, mock_requests):
mock_post, mock_get = mock_requests
_setup_mock_execution(mock_post, mock_get, num_images=2)
inputs = {"prompt": "c", "image_size": "square", "num_images": 2, "output_path": str(tmp_path / "c.jpeg")}
result = seedream_tool.execute(inputs)
assert result.cost_usd == pytest.approx(seedream_tool.estimate_cost(inputs))
def test_image_selector_routes_count_and_returns_distinct_artifacts(
self, seedream_tool, tmp_path, mock_requests, monkeypatch
):
from tools.graphics.image_selector import ImageSelector
mock_post, mock_get = mock_requests
_setup_mock_execution(mock_post, mock_get, num_images=2)
selector = ImageSelector()
monkeypatch.setattr(selector, "_providers", lambda: [seedream_tool])
result = selector.execute(
{
"prompt": "campaign artwork",
"preferred_provider": "bytedance",
"n": 2,
"output_path": str(tmp_path / "selected.png"),
}
)
assert result.success, result.error
assert result.data["selected_tool"] == "seedream_image"
assert len(set(result.artifacts)) == 2
assert {Path(path).name for path in result.artifacts} == {
"selected_1.png",
"selected_2.png",
}