import asyncio import time from types import SimpleNamespace from telethon import errors from bot.progress import ProgressReporter from bot.queue import Job, JobQueue def job(queue, user=1, name="a.mkv"): return Job(id=queue.new_id(), user_id=user, chat_id=user, msg_id=1, file_name=name, size=1, new_name=name) async def test_runs_jobs_fifo_with_single_worker(): order = [] gate = asyncio.Event() async def runner(j): order.append(("start", j.id)) await gate.wait() order.append(("end", j.id)) q = JobQueue(runner, workers=1) q.start() a, b, c = job(q), job(q), job(q) assert q.add(a) == 0 and q.add(b) == 1 and q.add(c) == 2 await asyncio.sleep(0.05) assert q.running == 1 and q.waiting == 2 assert q.position(a) == 0 and q.position(b) == 1 and q.position(c) == 2 gate.set() await asyncio.sleep(0.1) assert order == [("start", 1), ("end", 1), ("start", 2), ("end", 2), ("start", 3), ("end", 3)] assert q.jobs == {} # finished jobs are forgotten await q.stop() async def test_two_workers_run_in_parallel(): running = [] peak = 0 gate = asyncio.Event() async def runner(j): nonlocal peak running.append(j.id) peak = max(peak, len(running)) await gate.wait() running.remove(j.id) q = JobQueue(runner, workers=2) q.start() for _ in range(4): q.add(job(q)) await asyncio.sleep(0.05) assert q.running == 2 and peak == 2 gate.set() await asyncio.sleep(0.1) await q.stop() async def test_cancel_queued_job_skips_it_and_calls_back(): done, cancelled = [], [] gate = asyncio.Event() async def runner(j): done.append(j.id) await gate.wait() async def on_cancelled(j): cancelled.append(j.id) q = JobQueue(runner, workers=1, on_cancelled=on_cancelled) q.start() a, b, c = job(q), job(q), job(q) for j in (a, b, c): q.add(j) await asyncio.sleep(0.05) assert q.cancel(b.id) is True assert q.cancel(b.id) is False # already gone await asyncio.sleep(0.05) assert cancelled == [b.id] and b.state == "cancelled" gate.set() await asyncio.sleep(0.1) assert done == [a.id, c.id] await q.stop() async def test_cancel_running_job_sets_its_event_and_cancel_is_scoped_to_user(): seen = {} started = asyncio.Event() async def runner(j): started.set() await j.cancel.wait() seen["cancelled"] = True q = JobQueue(runner, workers=1) q.start() a = job(q, user=5) q.add(a) await started.wait() assert q.cancel(a.id, user_id=99) is False # someone else's job assert q.cancel(a.id, user_id=5) is True await asyncio.sleep(0.05) assert seen == {"cancelled": True} await q.stop() async def test_cancel_user_and_cancel_all(): gate = asyncio.Event() async def runner(j): await gate.wait() q = JobQueue(runner, workers=1) q.start() jobs = [job(q, user=1), job(q, user=1), job(q, user=2), job(q, user=2)] for j in jobs: q.add(j) await asyncio.sleep(0.05) assert q.cancel_user(1) == 2 assert [j.id for j in q.for_user(2)] == [jobs[2].id, jobs[3].id] assert q.cancel_all() == 2 gate.set() await q.stop() async def test_runner_crash_does_not_kill_the_worker(): ran = [] async def runner(j): ran.append(j.id) if j.id == 1: raise RuntimeError("boom") q = JobQueue(runner, workers=1) q.start() q.add(job(q)) q.add(job(q)) await asyncio.sleep(0.1) assert ran == [1, 2] await q.stop() def ordered_job(queue, order, user=1): return Job(id=queue.new_id(), user_id=user, chat_id=user, msg_id=order, file_name=f"{order}.mkv", size=1, new_name=f"{order}.mkv", order=order) async def idle_queue(runner, **kw): """A started queue whose worker is already waiting for work, like the real bot between files.""" q = JobQueue(runner, workers=1, **kw) q.start() await asyncio.sleep(0.02) return q async def test_a_batch_queued_out_of_order_runs_in_message_order(): ran = [] async def runner(j): ran.append(j.order) q = await idle_queue(runner, settle=0.1) for order in (1146, 1147, 1145): # the order in which three concurrent handlers happened to finish q.add(ordered_job(q, order)) await asyncio.sleep(0.01) await asyncio.sleep(0.4) assert ran == [1145, 1146, 1147] await q.stop() async def test_without_a_settle_time_the_first_queued_job_starts_at_once(): ran = [] async def runner(j): ran.append(j.order) q = await idle_queue(runner) for order in (1146, 1147, 1145): q.add(ordered_job(q, order)) await asyncio.sleep(0.01) await asyncio.sleep(0.1) assert ran[0] == 1146 # which is why the bot uses a settle time await q.stop() async def test_jobs_queued_while_busy_are_picked_oldest_message_first(): ran, gate = [], asyncio.Event() async def runner(j): ran.append(j.order) if j.order == 1: await gate.wait() q = await idle_queue(runner) q.add(ordered_job(q, 1)) await asyncio.sleep(0.02) for order in (30, 20, 25): q.add(ordered_job(q, order)) gate.set() await asyncio.sleep(0.1) assert ran == [1, 20, 25, 30] await q.stop() async def test_an_idle_worker_waits_for_the_settle_time_before_starting(): ran = [] async def runner(j): ran.append(j.id) q = await idle_queue(runner, settle=0.3) q.add(job(q)) await asyncio.sleep(0.1) assert ran == [] and q.waiting == 1 await asyncio.sleep(0.5) assert ran == [1] await q.stop() async def test_every_new_job_extends_the_wait_so_a_slow_batch_is_still_grouped(): ran = [] async def runner(j): ran.append(j.order) q = await idle_queue(runner, settle=0.3) for order in (3, 2, 1): # each arrives 0.2 s after the previous one: more than one settle period in total q.add(ordered_job(q, order)) await asyncio.sleep(0.2) await asyncio.sleep(0.6) assert ran == [1, 2, 3] await q.stop() async def test_a_steady_stream_cannot_hold_the_worker_back_forever(): started_at, last_add = [], [] async def runner(j): started_at.append(time.monotonic()) q = await idle_queue(runner, settle=0.15) # capped at SETTLE_CAP * 0.15 s for order in range(30, 0, -1): q.add(ordered_job(q, order)) last_add.append(time.monotonic()) await asyncio.sleep(0.05) assert started_at and started_at[0] < last_add[-1] - 0.3 # it started well before the stream stopped await q.stop() async def test_position_and_listing_follow_message_order(): async def runner(j): pass q = JobQueue(runner, workers=1) # never started: everything stays queued a, b, c = ordered_job(q, 30), ordered_job(q, 10), ordered_job(q, 20) for j in (a, b, c): q.add(j) assert [q.position(j) for j in (b, c, a)] == [1, 2, 3] assert [j.order for j in q.for_user(1)] == [10, 20, 30] async def test_cancelling_during_the_settle_time_skips_the_job(): ran = [] async def runner(j): ran.append(j.id) q = await idle_queue(runner, settle=0.2) a, b = job(q), job(q) q.add(a) q.add(b) await asyncio.sleep(0.05) assert q.cancel(a.id) await asyncio.sleep(0.6) assert ran == [b.id] await q.stop() # --------------------------------------------------------------------------- # progress reporter # --------------------------------------------------------------------------- class FakeClient: def __init__(self): self.edits = [] self.error = None async def edit_message(self, message, text, buttons=None, parse_mode=None): if self.error: raise self.error self.edits.append((text, buttons)) async def test_reporter_throttles_and_shows_stage_and_cancel_button(): client = FakeClient() r = ProgressReporter(client, SimpleNamespace(id=1), 7, "My .mkv", interval=60) r.set_stage("⬇️ Downloading", 1000) await asyncio.sleep(0.02) assert len(client.edits) == 1 text, buttons = client.edits[0] assert "Downloading" in text and "My <File>.mkv" in text assert buttons[0][0].type.data == b"job:x:7" for cur in range(1, 500, 50): r(cur, 1000) # within the throttle window: no more edits await asyncio.sleep(0.02) assert len(client.edits) == 1 r(1000, 1000) # completion always shows await asyncio.sleep(0.02) assert len(client.edits) == 2 and "100.0%" in client.edits[-1][0] async def test_stage_change_is_shown_even_while_an_edit_is_in_flight(): class SlowClient(FakeClient): async def edit_message(self, message, text, buttons=None, parse_mode=None): await asyncio.sleep(0.05) await super().edit_message(message, text, buttons=buttons, parse_mode=parse_mode) client = SlowClient() r = ProgressReporter(client, None, 1, "x", interval=0) r.set_stage("⬇️ Downloading", 100) r(100, 100) # dropped: the first edit is still running r.set_stage("⚙️ Applying metadata", 0) # must not be dropped await r.finish("✅ Done") stages = [t for t, _b in client.edits] assert "Downloading" in stages[0] and "Applying metadata" in stages[1] and stages[-1] == "✅ Done" assert len(stages) == 3 async def test_reporter_render_numbers(): r = ProgressReporter(FakeClient(), None, 1, "x", interval=4) r.stage = "⬆️ Uploading" r._speed = 1024 * 1024 text = r.render(512 * 1024 * 1024, 1024 * 1024 * 1024) assert "50.0%" in text and "512.00 MB / 1.00 GB" in text and "1.00 MB/s" in text and "ETA 8:32" in text assert "Working" in r.render(0, 0) async def test_reporter_survives_flood_wait_and_message_errors(): client = FakeClient() r = ProgressReporter(client, None, 1, "x", interval=0) client.error = errors.FloodWaitError(request=None, capture=3) r.set_stage("a", 10) await asyncio.sleep(0.02) assert r._muted_until > 0 # backed off instead of crashing client.error = errors.MessageNotModifiedError(request=None) r._muted_until = 0 r(5, 10) await asyncio.sleep(0.02) client.error = RuntimeError("deleted") r(10, 10) await asyncio.sleep(0.02) async def test_reporter_finish_removes_buttons(): client = FakeClient() r = ProgressReporter(client, None, 1, "x") await r.finish("✅ Done") assert client.edits[-1] == ("✅ Done", None)