75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes._execution_helpers import advance_progress, frame_progress, raise_if_interrupted
|
|
|
|
|
|
class FakeInterrupt(BaseException):
|
|
"""Mirrors ComfyUI's InterruptProcessingException: not an Exception."""
|
|
|
|
|
|
class RecordingBar:
|
|
def __init__(self):
|
|
self.updates = []
|
|
|
|
def update_absolute(self, value, total=None, preview=None):
|
|
self.updates.append((value, total, preview))
|
|
|
|
|
|
class BrokenBar:
|
|
def update_absolute(self, *_args, **_kwargs):
|
|
raise RuntimeError("the websocket went away")
|
|
|
|
|
|
class InterruptingBar:
|
|
"""ComfyUI's progress hook raises its interrupt straight through here."""
|
|
|
|
def update_absolute(self, *_args, **_kwargs):
|
|
raise FakeInterrupt()
|
|
|
|
|
|
class FrameProgressTests(unittest.TestCase):
|
|
def test_an_unknown_total_tracks_nothing(self):
|
|
self.assertIsNone(frame_progress(0))
|
|
self.assertIsNone(frame_progress(-7))
|
|
|
|
def test_a_known_total_either_tracks_or_fails_soft(self):
|
|
# A bar only exists when ComfyUI is importable; either way no throw.
|
|
bar = frame_progress(10)
|
|
self.assertTrue(bar is None or hasattr(bar, "update_absolute"))
|
|
|
|
|
|
class AdvanceProgressTests(unittest.TestCase):
|
|
def test_an_absent_bar_is_a_no_op(self):
|
|
advance_progress(None, 3, 10)
|
|
|
|
def test_updates_are_absolute_and_carry_no_preview(self):
|
|
bar = RecordingBar()
|
|
advance_progress(bar, 3, 10)
|
|
advance_progress(bar, 4, 10)
|
|
self.assertEqual(bar.updates, [(3, 10, None), (4, 10, None)])
|
|
|
|
def test_a_broken_bar_never_breaks_the_operation(self):
|
|
advance_progress(BrokenBar(), 1, 10)
|
|
|
|
def test_an_interrupt_raised_by_the_bar_still_gets_out(self):
|
|
with self.assertRaises(FakeInterrupt):
|
|
advance_progress(InterruptingBar(), 1, 10)
|
|
|
|
|
|
class InterruptSeamTests(unittest.TestCase):
|
|
def test_an_uninterrupted_queue_lets_the_loop_continue(self):
|
|
self.assertIsNone(raise_if_interrupted())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|