mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-09-26 07:57:50 -05:00
97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
import types
|
|
from unittest.mock import Mock
|
|
|
|
import comfy_execution.progress as progress
|
|
from comfy_execution.progress import ProgressRegistry, WebUIProgressHandler
|
|
|
|
|
|
class Clock:
|
|
def __init__(self):
|
|
self.now = 100.0
|
|
|
|
def advance(self, seconds):
|
|
self.now += seconds
|
|
|
|
def perf_counter(self):
|
|
return self.now
|
|
|
|
|
|
def make_registry(monkeypatch):
|
|
clock = Clock()
|
|
monkeypatch.setattr(progress, "time", types.SimpleNamespace(perf_counter=clock.perf_counter))
|
|
dynprompt = Mock()
|
|
dynprompt.get_display_node_id.side_effect = lambda n: n
|
|
dynprompt.get_parent_node_id.return_value = None
|
|
dynprompt.get_real_node_id.side_effect = lambda n: n
|
|
dynprompt.get_node.return_value = {"class_type": "KSampler"}
|
|
return ProgressRegistry(prompt_id="p", dynprompt=dynprompt), clock
|
|
|
|
|
|
def test_seconds_sum_every_stretch_of_an_activity(monkeypatch):
|
|
registry, clock = make_registry(monkeypatch)
|
|
registry.start_progress("1")
|
|
clock.advance(1.0)
|
|
registry.set_activity("1", "loading")
|
|
clock.advance(2.0)
|
|
registry.set_activity("1", None)
|
|
clock.advance(0.5)
|
|
registry.set_activity("1", "loading")
|
|
clock.advance(3.0)
|
|
registry.set_activity("1", None)
|
|
clock.advance(0.25)
|
|
registry.finish_progress("1")
|
|
|
|
assert registry.nodes["1"]["seconds"] == {"loading": 5.0, "total": 6.75}
|
|
assert "activity" not in registry.nodes["1"]
|
|
|
|
|
|
def test_switching_activities_closes_the_previous_stretch(monkeypatch):
|
|
registry, clock = make_registry(monkeypatch)
|
|
registry.start_progress("1")
|
|
registry.set_activity("1", "loading")
|
|
clock.advance(1.0)
|
|
registry.set_activity("1", "downloading")
|
|
clock.advance(4.0)
|
|
registry.finish_progress("1")
|
|
|
|
assert registry.nodes["1"]["seconds"] == {"loading": 1.0, "downloading": 4.0, "total": 5.0}
|
|
|
|
|
|
def test_cached_node_reports_no_seconds(monkeypatch):
|
|
registry, _ = make_registry(monkeypatch)
|
|
registry.finish_progress("1")
|
|
|
|
assert "seconds" not in registry.nodes["1"]
|
|
|
|
|
|
def test_two_nodes_are_timed_separately(monkeypatch):
|
|
registry, clock = make_registry(monkeypatch)
|
|
registry.start_progress("1")
|
|
registry.set_activity("1", "downloading")
|
|
clock.advance(1.0)
|
|
registry.start_progress("2")
|
|
registry.set_activity("2", "loading")
|
|
clock.advance(2.0)
|
|
registry.set_activity("1", None)
|
|
registry.finish_progress("1")
|
|
clock.advance(1.0)
|
|
registry.finish_progress("2")
|
|
|
|
assert registry.nodes["1"]["seconds"] == {"downloading": 3.0, "total": 3.0}
|
|
assert registry.nodes["2"]["seconds"] == {"loading": 3.0, "total": 3.0}
|
|
|
|
|
|
def test_seconds_are_sent_only_once_the_node_finishes(monkeypatch):
|
|
registry, clock = make_registry(monkeypatch)
|
|
server = Mock(client_id="c")
|
|
handler = WebUIProgressHandler(server)
|
|
handler.set_registry(registry)
|
|
registry.register_handler(handler)
|
|
|
|
registry.start_progress("1")
|
|
assert "seconds" not in server.send_sync.call_args[0][1]["nodes"]["1"]
|
|
|
|
clock.advance(2.0)
|
|
registry.finish_progress("1")
|
|
assert server.send_sync.call_args[0][1]["nodes"]["1"]["seconds"] == {"total": 2.0}
|