pplate/pi_mcps
1
1
Fork 0
Code Issues 9 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

docs(mcp-image-gen): add USAGE.md and expand tests to 19

Browse Source
This commit is contained in:
pplate
2026-04-04 12:16:03 +02:00
parent b0ce5c55ed
commit 8cbeb6571b
2 changed files with 838 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mcp/mcp-image-gen/tests/test_server.py
+250 -5
View File
@@ -28,7 +28,6 @@ COMFYUI_BASE = "http://test-comfyui:8188"
# build_flux_workflow — pure function, no mocking needed
# ---------------------------------------------------------------------------
def test_build_flux_workflow_structure():
"""Verify build_flux_workflow returns a dict with correct node types."""
wf = build_flux_workflow(
@@ -103,7 +102,6 @@ def test_random_seed_generated():
# list_available_models
# ---------------------------------------------------------------------------
@respx.mock
@pytest.mark.asyncio
async def test_list_available_models():
@@ -146,7 +144,6 @@ async def test_list_available_models_comfyui_offline():
# get_generation_status
# ---------------------------------------------------------------------------
@respx.mock
@pytest.mark.asyncio
async def test_get_generation_status_pending(queue_with_pending):
@@ -191,7 +188,6 @@ async def test_get_generation_status_complete(queue_empty, mock_history_response
# get_output_directory
# ---------------------------------------------------------------------------
def test_get_output_directory_default(monkeypatch):
"""No IMAGE_OUTPUT_DIR env var → returns expanded ~/Pictures/mcp-generated."""
monkeypatch.delenv("IMAGE_OUTPUT_DIR", raising=False)
@@ -216,7 +212,6 @@ def test_get_output_directory_custom(monkeypatch, tmp_path):
# generate_image
# ---------------------------------------------------------------------------
@respx.mock
@pytest.mark.asyncio
async def test_generate_image_success(
@@ -300,3 +295,253 @@ async def test_generate_image_timeout(monkeypatch, queue_with_pending):
assert len(result) == 1
assert "timed out" in result[0].text.lower()
assert "test-uuid-1234" in result[0].text
@respx.mock
@pytest.mark.asyncio
async def test_generate_image_empty_prompt(tmp_path, sample_image_bytes, mock_history_response, queue_empty, monkeypatch):
"""Empty prompt → workflow has empty text in positive node, but generation succeeds."""
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path))
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
return_value=httpx.Response(200, json={"prompt_id": "test-empty-uuid"})
)
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
return_value=httpx.Response(200, json=queue_empty)
)
mock_history_empty = {
"test-empty-uuid": {
"outputs": {
"9": {
"images": [
{
"filename": "mcp-image-gen_00001_.png",
"subfolder": "",
"type": "output",
}
]
}
},
"status": {"completed": True},
}
}
respx.get(f"{COMFYUI_BASE}/api/history/test-empty-uuid").mock(
return_value=httpx.Response(200, json=mock_history_empty)
)
respx.get(f"{COMFYUI_BASE}/api/view").mock(
return_value=httpx.Response(200, content=sample_image_bytes)
)
result = await generate_image(prompt="", output_dir=str(tmp_path))
assert len(result) == 2
text_content = result[0]
image_content = result[1]
assert "Generated:" in text_content.text
assert str(tmp_path) in text_content.text
# Verify workflow was built with empty prompt (indirectly via success)
assert image_content.mimeType == "image/png"
@respx.mock
@pytest.mark.asyncio
async def test_generate_image_long_prompt(tmp_path, sample_image_bytes, mock_history_response, queue_empty, monkeypatch):
"""Very long prompt → passed as-is to workflow without truncation."""
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path))
long_prompt = "a " + "very long descriptive prompt " * 50 # ~500 chars
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
return_value=httpx.Response(200, json={"prompt_id": "test-long-uuid"})
)
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
return_value=httpx.Response(200, json=queue_empty)
)
mock_history_long = {
"test-long-uuid": {
"outputs": {
"9": {
"images": [
{
"filename": "mcp-image-gen_00001_.png",
"subfolder": "",
"type": "output",
}
]
}
},
"status": {"completed": True},
}
}
respx.get(f"{COMFYUI_BASE}/api/history/test-long-uuid").mock(
return_value=httpx.Response(200, json=mock_history_long)
)
respx.get(f"{COMFYUI_BASE}/api/view").mock(
return_value=httpx.Response(200, content=sample_image_bytes)
)
result = await generate_image(prompt=long_prompt, output_dir=str(tmp_path))
assert len(result) == 2
# Success implies long prompt was accepted (ComfyUI handles it)
saved_files = list(tmp_path.glob("*.png"))
assert len(saved_files) == 1
@respx.mock
@pytest.mark.asyncio
async def test_generate_image_invalid_model(tmp_path, monkeypatch):
"""Invalid model → ComfyUI /prompt returns 500 or 404, tool returns error TextContent."""
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path))
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
return_value=httpx.Response(404, json={"error": "Model not found"})
)
result = await generate_image(
prompt="a cat",
model="nonexistent-model.safetensors",
output_dir=str(tmp_path)
)
assert len(result) == 1
assert "404" in result[0].text
assert "Model not found" in result[0].text
# No file saved
saved_files = list(tmp_path.glob("*.png"))
assert len(saved_files) == 0
@respx.mock
@pytest.mark.asyncio
async def test_generate_image_custom_output_dir(tmp_path, sample_image_bytes, mock_history_response, queue_empty, monkeypatch):
"""Custom output_dir → image saved there, path reflects it."""
custom_dir = tmp_path / "custom"
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path)) # Base for default
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
return_value=httpx.Response(200, json={"prompt_id": "test-custom-uuid"})
)
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
return_value=httpx.Response(200, json=queue_empty)
)
mock_history_custom = {
"test-custom-uuid": {
"outputs": {
"9": {
"images": [
{
"filename": "mcp-image-gen_00001_.png",
"subfolder": "",
"type": "output",
}
]
}
},
"status": {"completed": True},
}
}
respx.get(f"{COMFYUI_BASE}/api/history/test-custom-uuid").mock(
return_value=httpx.Response(200, json=mock_history_custom)
)
respx.get(f"{COMFYUI_BASE}/api/view").mock(
return_value=httpx.Response(200, content=sample_image_bytes)
)
result = await generate_image(
prompt="a dog",
output_dir=str(custom_dir),
)
assert len(result) == 2
text_content = result[0]
assert str(custom_dir) in text_content.text
# Directory was created
assert custom_dir.exists()
saved_files = list(custom_dir.glob("*.png"))
assert len(saved_files) == 1
@respx.mock
@pytest.mark.asyncio
async def test_generate_image_random_seed_variance(tmp_path, sample_image_bytes, mock_history_response, queue_empty, monkeypatch):
"""seed=-1 → different actual_seed each call, reflected in filename."""
monkeypatch.setattr(server, "IMAGE_OUTPUT_DIR", str(tmp_path))
# First generation
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
return_value=httpx.Response(200, json={"prompt_id": "seed1-uuid"})
)
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
return_value=httpx.Response(200, json=queue_empty)
)
mock_history_seed1 = {
"seed1-uuid": {
"outputs": {
"9": {
"images": [
{
"filename": "mcp-image-gen_00001_.png",
"subfolder": "",
"type": "output",
}
]
}
},
"status": {"completed": True},
}
}
respx.get(f"{COMFYUI_BASE}/api/history/seed1-uuid").mock(
return_value=httpx.Response(200, json=mock_history_seed1)
)
respx.get(f"{COMFYUI_BASE}/api/view").mock(
return_value=httpx.Response(200, content=sample_image_bytes)
)
result1 = await generate_image(prompt="cat", seed=-1, output_dir=str(tmp_path))
seed1 = [line for line in result1[0].text.split("\n") if "Seed:" in line][0].split(": ")[1]
filename1 = Path(result1[0].text.split("Generated: ")[1].split("\n")[0]).name
assert "Seed:" in result1[0].text
assert int(seed1) != 0 # Not default
# Reset mocks for second call
respx.reset()
respx.post(f"{COMFYUI_BASE}/api/prompt").mock(
return_value=httpx.Response(200, json={"prompt_id": "seed2-uuid"})
)
respx.get(f"{COMFYUI_BASE}/api/queue").mock(
return_value=httpx.Response(200, json=queue_empty)
)
mock_history_seed2 = {
"seed2-uuid": {
"outputs": {
"9": {
"images": [
{
"filename": "mcp-image-gen_00002_.png",
"subfolder": "",
"type": "output",
}
]
}
},
"status": {"completed": True},
}
}
respx.get(f"{COMFYUI_BASE}/api/history/seed2-uuid").mock(
return_value=httpx.Response(200, json=mock_history_seed2)
)
respx.get(f"{COMFYUI_BASE}/api/view").mock(
return_value=httpx.Response(200, content=sample_image_bytes)
)
result2 = await generate_image(prompt="cat", seed=-1, output_dir=str(tmp_path))
seed2 = [line for line in result2[0].text.split("\n") if "Seed:" in line][0].split(": ")[1]
filename2 = Path(result2[0].text.split("Generated: ")[1].split("\n")[0]).name
# Different seeds and filenames
assert seed1 != seed2
assert filename1 != filename2
# Both saved
saved_files = list(tmp_path.glob("*.png"))
assert len(saved_files) == 2