mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 12:28:17 +08:00
118 lines
4.2 KiB
Python
118 lines
4.2 KiB
Python
import execution
|
|
|
|
|
|
class TestServer:
|
|
def queue_updated(self):
|
|
pass
|
|
|
|
|
|
def test_cooperative_cancel_is_prompt_scoped_without_node_class_checks(monkeypatch):
|
|
interrupted = []
|
|
monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
|
|
queue = execution.PromptQueue(TestServer())
|
|
prompt = {"1": {"class_type": "UnrelatedNode", "inputs": {}}}
|
|
queue.put((0, "first", prompt, {}, [], {}))
|
|
queue.put((1, "second", prompt, {}, [], {}))
|
|
_, first_id = queue.get(timeout=0)
|
|
_, second_id = queue.get(timeout=0)
|
|
|
|
queue.set_cooperative(True)
|
|
assert queue.interrupt_if_running("first")
|
|
assert queue.is_cancelled("first")
|
|
assert not queue.is_cancelled("second")
|
|
assert interrupted == []
|
|
|
|
queue.task_done(first_id, {}, None)
|
|
assert not queue.is_cancelled("first")
|
|
queue.task_done(second_id, {}, None)
|
|
|
|
|
|
def test_single_cooperative_cancel_also_wakes_synchronous_execution(monkeypatch):
|
|
interrupted = []
|
|
monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
|
|
queue = execution.PromptQueue(TestServer())
|
|
queue.put((0, "only", {}, {}, [], {}))
|
|
_, item_id = queue.get(timeout=0)
|
|
queue.set_cooperative(True)
|
|
|
|
assert queue.interrupt_if_running("only")
|
|
assert queue.is_cancelled("only")
|
|
assert queue.is_cooperative_draining()
|
|
assert interrupted == [True]
|
|
|
|
queue.task_done(item_id, {}, None)
|
|
queue.finish_cooperative_drain()
|
|
|
|
|
|
def test_completed_result_wins_if_cancel_arrives_after_compute(monkeypatch):
|
|
monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: None)
|
|
queue = execution.PromptQueue(TestServer())
|
|
queue.put((0, "completed", {}, {}, [], {}))
|
|
_, item_id = queue.get(timeout=0)
|
|
queue.set_cooperative(True)
|
|
|
|
assert queue.interrupt_if_running("completed")
|
|
status = execution.PromptQueue.ExecutionStatus("success", True, [])
|
|
queue.task_done(item_id, {"outputs": {"node": "result"}}, status)
|
|
|
|
assert not queue.is_cancelled("completed")
|
|
history = queue.get_history("completed")["completed"]
|
|
assert history["status"]["status_str"] == "success"
|
|
assert history["outputs"] == {"node": "result"}
|
|
|
|
|
|
def test_global_cooperative_interrupt_cancels_all_and_blocks_admission_until_drain(monkeypatch):
|
|
interrupted = []
|
|
monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
|
|
queue = execution.PromptQueue(TestServer())
|
|
for number, prompt_id in enumerate(("first", "second", "next")):
|
|
queue.put((number, prompt_id, {}, {}, [], {}))
|
|
_, first_id = queue.get(timeout=0)
|
|
_, second_id = queue.get(timeout=0)
|
|
queue.set_cooperative(True)
|
|
|
|
queue.interrupt_all_running()
|
|
|
|
assert queue.is_cancelled("first")
|
|
assert queue.is_cancelled("second")
|
|
assert interrupted == [True]
|
|
assert queue.get_if(lambda item: True) is None
|
|
|
|
queue.task_done(first_id, {}, None)
|
|
queue.task_done(second_id, {}, None)
|
|
queue.finish_cooperative_drain()
|
|
item, next_id = queue.get_if(lambda item: True)
|
|
assert item[1] == "next"
|
|
queue.task_done(next_id, {}, None)
|
|
|
|
|
|
def test_legacy_cancel_uses_global_interrupt(monkeypatch):
|
|
interrupted = []
|
|
monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
|
|
queue = execution.PromptQueue(TestServer())
|
|
queue.put((0, "prompt", {}, {}, [], {}))
|
|
queue.get(timeout=0)
|
|
|
|
assert queue.interrupt_if_running("prompt")
|
|
assert interrupted == [True]
|
|
assert not queue.is_cancelled("prompt")
|
|
|
|
|
|
def test_targeted_cancel_does_not_interrupt_different_serial_prompt(monkeypatch):
|
|
interrupted = []
|
|
monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
|
|
queue = execution.PromptQueue(TestServer())
|
|
queue.put((0, "running", {}, {}, [], {}))
|
|
queue.get(timeout=0)
|
|
|
|
assert not queue.interrupt_if_running("different")
|
|
assert interrupted == []
|
|
|
|
|
|
def test_get_if_leaves_incompatible_head_queued():
|
|
queue = execution.PromptQueue(TestServer())
|
|
queue.put((0, "first", {}, {}, [], {}))
|
|
|
|
assert queue.get_if(lambda item: item[1] == "other") is None
|
|
assert queue.get(timeout=0)[0][1] == "first"
|