Spaces:
Running on Zero
Running on Zero
File size: 2,997 Bytes
cf8e3da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import fallback_progress # noqa: E402
class FakeClock:
def __init__(self):
self.now = 0.0
def __call__(self):
return self.now
def sleep(self, seconds):
self.now += seconds
def run(entry, levels, finish_after=None, clock=None):
clock = clock or FakeClock()
calls, reads = [], iter(levels)
def read(event_id):
if finish_after is not None and clock.now >= finish_after:
entry.status = "complete"
return next(reads, levels[-1] if levels else None)
fallback_progress.mirror_progress(
entry, lambda pct, desc=None: calls.append((pct, desc)), read=read, sleep=clock.sleep, clock=clock,
)
return calls, clock
class MirrorProgress(unittest.TestCase):
def entry(self, **fields):
return SimpleNamespace(**{"status": "running", "cancelled": False, "progress_event": "evt", **fields})
def test_forwards_each_new_level_once_until_the_run_completes(self):
levels = [None, None, (0.02, "Preparing image"), (0.02, "Preparing image"), (0.15, "Generating 3D structure")]
calls, _ = run(self.entry(), levels, finish_after=2.0)
self.assertEqual(calls, [(0.02, "Preparing image"), (0.15, "Generating 3D structure")])
def test_stays_silent_while_the_pro_run_has_not_reached_the_gpu(self):
calls, clock = run(self.entry(), [None], finish_after=5.0)
self.assertEqual(calls, [])
self.assertLess(clock.now, 6.0)
def test_finished_or_cancelled_entries_return_immediately(self):
for entry in (self.entry(status="complete"), self.entry(status="failed"), self.entry(cancelled=True)):
calls, clock = run(entry, [(0.5, "x")])
self.assertEqual((calls, clock.now), ([], 0.0))
def test_gives_up_after_the_timeout(self):
_, clock = run(self.entry(), [None])
self.assertGreaterEqual(clock.now, fallback_progress.TIMEOUT_SECONDS)
self.assertLess(clock.now, fallback_progress.TIMEOUT_SECONDS + 1)
def test_event_progress_reads_the_latest_unit_of_the_matching_event(self):
unit = SimpleNamespace(progress=0.45, desc="Rendering preview")
event = SimpleNamespace(_id="evt", progress=SimpleNamespace(progress_data=[unit]))
other = SimpleNamespace(_id="other", progress=None)
blocks = SimpleNamespace(_queue=SimpleNamespace(active_jobs=[None, [other, event]]))
token = fallback_progress.LocalContext.blocks.set(blocks)
try:
self.assertEqual(fallback_progress.event_progress("evt"), (0.45, "Rendering preview"))
self.assertIsNone(fallback_progress.event_progress("missing"))
self.assertIsNone(fallback_progress.event_progress(None))
finally:
fallback_progress.LocalContext.blocks.reset(token)
if __name__ == "__main__":
unittest.main()
|