mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-05 18:05:08 +08:00
afb8cbbe3e
Amp-Thread-ID: https://ampcode.com/threads/T-019fca67-53e6-76d6-8561-0bb1b07473d1 Co-authored-by: Amp <amp@ampcode.com>
107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
import asyncio
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from comfy.cli_args import args
|
|
|
|
if not torch.cuda.is_available():
|
|
args.cpu = True
|
|
|
|
from comfy_api_nodes.apis.comfy_cloud import (
|
|
ComfyCloudGenerateRequest,
|
|
ComfyCloudGenerateResponse,
|
|
ComfyCloudStatusResponse,
|
|
ComfyCloudWorkflowInputs,
|
|
)
|
|
from comfy_api_nodes import nodes_comfy_cloud
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("node", "workflow", "returns_video", "requires_image"),
|
|
[
|
|
(nodes_comfy_cloud.ComfyCloudTextToImageNode, "text-to-image", False, False),
|
|
(nodes_comfy_cloud.ComfyCloudTextToVideoNode, "text-to-video", True, False),
|
|
(nodes_comfy_cloud.ComfyCloudImageToVideoNode, "image-to-video", True, True),
|
|
(nodes_comfy_cloud.ComfyCloudImageEditNode, "image-edit", False, True),
|
|
],
|
|
)
|
|
def test_workflow_submission_polling_and_download(monkeypatch, node, workflow, returns_video, requires_image):
|
|
sync = AsyncMock(
|
|
return_value=ComfyCloudGenerateResponse(
|
|
task_id="task-1",
|
|
status="queued",
|
|
polling_url="/proxy/comfy-cloud/workflow/tasks/task-1",
|
|
cancel_url="/proxy/comfy-cloud/workflow/tasks/task-1/cancel",
|
|
)
|
|
)
|
|
poll = AsyncMock(
|
|
return_value=ComfyCloudStatusResponse(
|
|
task_id="task-1",
|
|
status="completed",
|
|
progress=100,
|
|
output_url="https://example.com/output",
|
|
)
|
|
)
|
|
upload = AsyncMock(return_value="https://example.com/input.png")
|
|
image_download = AsyncMock(return_value="image-output")
|
|
video_download = AsyncMock(return_value="video-output")
|
|
monkeypatch.setattr(nodes_comfy_cloud, "sync_op", sync)
|
|
monkeypatch.setattr(nodes_comfy_cloud, "poll_op", poll)
|
|
monkeypatch.setattr(nodes_comfy_cloud, "upload_image_to_comfyapi", upload)
|
|
monkeypatch.setattr(nodes_comfy_cloud, "download_url_to_image_tensor", image_download)
|
|
monkeypatch.setattr(nodes_comfy_cloud, "download_url_to_video_output", video_download)
|
|
monkeypatch.setattr(nodes_comfy_cloud, "get_number_of_images", lambda image: 1)
|
|
|
|
image = object() if requires_image else None
|
|
output = asyncio.run(node.execute("A tiny fennec fox", image))
|
|
|
|
endpoint = sync.call_args.args[1]
|
|
request = sync.call_args.kwargs["data"]
|
|
assert endpoint.path == "/proxy/comfy-cloud/workflow/generate"
|
|
assert endpoint.method == "POST"
|
|
assert request == ComfyCloudGenerateRequest(
|
|
workflow=workflow,
|
|
inputs=ComfyCloudWorkflowInputs(
|
|
prompt="A tiny fennec fox",
|
|
image_url="https://example.com/input.png" if requires_image else None,
|
|
),
|
|
)
|
|
assert upload.await_count == int(requires_image)
|
|
|
|
poll_endpoint = poll.call_args.args[1]
|
|
cancel_endpoint = poll.call_args.kwargs["cancel_endpoint"]
|
|
assert poll_endpoint.path == "/proxy/comfy-cloud/workflow/tasks/task-1"
|
|
assert cancel_endpoint.path == "/proxy/comfy-cloud/workflow/tasks/task-1/cancel"
|
|
assert cancel_endpoint.method == "POST"
|
|
assert output[0] == ("video-output" if returns_video else "image-output")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"node",
|
|
[nodes_comfy_cloud.ComfyCloudImageToVideoNode, nodes_comfy_cloud.ComfyCloudImageEditNode],
|
|
)
|
|
def test_image_workflows_reject_batches(monkeypatch, node):
|
|
upload = AsyncMock()
|
|
monkeypatch.setattr(nodes_comfy_cloud, "get_number_of_images", lambda image: 2)
|
|
monkeypatch.setattr(nodes_comfy_cloud, "upload_image_to_comfyapi", upload)
|
|
|
|
with pytest.raises(ValueError, match="Exactly one input image"):
|
|
asyncio.run(node.execute("Animate this", object()))
|
|
upload.assert_not_awaited()
|
|
|
|
|
|
def test_contract_omits_optional_status_fields():
|
|
request = ComfyCloudGenerateRequest(
|
|
workflow="text-to-image",
|
|
inputs=ComfyCloudWorkflowInputs(prompt="A lighthouse"),
|
|
)
|
|
status = ComfyCloudStatusResponse(task_id="task-1", status="queued")
|
|
|
|
assert request.model_dump(exclude_none=True) == {
|
|
"workflow": "text-to-image",
|
|
"inputs": {"prompt": "A lighthouse"},
|
|
}
|
|
assert status.model_dump(exclude_none=True) == {"task_id": "task-1", "status": "queued"}
|