Spaces:
Running on Zero
Running on Zero
| 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() | |