import asyncio import hashlib import shutil import pytest from conftest import ( CHANNEL_ID, EXAMPLE_CAPTION, FakeDrive, FakeSource, RecordingNotifier, SimulatedCrash, build_queue, make_config, run, video_bytes, ) from database import Database from drive_client import SOURCE_PROPERTY from errors import PermanentError, TransientError from models import State from progress import ProgressReporter def caption(title, topic, batch="English ( SSC MAINS )", index="001"): return f"Index: {index}\nTitle: {title}\nTopic: {topic}\nBatch: {batch}\nExtracted By: @someone" async def sync_and_run(queue): await queue.sync() queue.start() await queue.wait() def key(message_id): return (CHANNEL_ID, message_id) def drive_names(drive, folder_name=None): names = [] for f in drive.uploaded(): parent = drive.files[f.parents[0]].name if folder_name is None or parent == folder_name: names.append(f.name) return sorted(names) # -- end to end --------------------------------------------------------------------------- def test_full_pipeline_topics_and_metadata(config, db, source, drive): source.add(1, EXAMPLE_CAPTION) source.add(2, caption("Class 02 Aman sir.mp4", "Misc• Grammar Practice")) source.add(3, caption("Class 01 Aman sir.mp4", "Error Detection")) source.add_text(4, "just an announcement") queue = build_queue(config, db, source, drive) run(sync_and_run(queue)) assert db.counts()["COMPLETED"] == 3 assert drive_names(drive, "Misc• Grammar Practice") == ["Class 01 Aman sir.mp4", "Class 02 Aman sir.mp4"] assert drive_names(drive, "Error Detection") == ["Class 01 Aman sir.mp4"] assert drive.created_folders == 2 # topic folder reused for the second video rec = db.get(key(1)) assert rec.indexes == ["001", "002"] assert rec.batch == "English ( SSC MAINS )" and rec.extracted_by == "@PUTINxINDIA" assert rec.final_name == "Class 01 Aman sir.mp4" uploaded = drive.files[rec.drive_file_id] assert uploaded.app_properties[SOURCE_PROPERTY] == f"{CHANNEL_ID}:1" assert uploaded.size == rec.local_size == rec.expected_size assert uploaded.md5 == rec.md5 assert rec.sha256 and rec.completed_at assert rec.download_status == "done" and rec.upload_status == "done" # temporary files are removed only after completion assert list(config.download_dir.iterdir()) == [] assert queue.stats.uploaded == 3 assert db.get_int("last_completed_message_id") == 3 assert db.get_int("last_discovered_message_id") == 4 def test_missing_topic_and_title_use_fallbacks(config, db, source, drive): source.add(1, "Title: Lonely.mp4") source.add(2, "Topic: Grammar", file_name="original_upload.mp4") source.add(3, "") queue = build_queue(config, db, source, drive) run(sync_and_run(queue)) assert drive_names(drive, "Uncategorized") == ["Lonely.mp4", "telegram_3.mp4"] assert drive_names(drive, "Grammar") == ["original_upload.mp4"] assert any("Title missing" in w for w in db.get(key(2)).warnings) def test_batch_subfolders(tmp_path, source, drive): config = make_config(tmp_path, BATCH_SUBFOLDERS_ENABLED="true") db = Database.from_config(config) source.add(1, EXAMPLE_CAPTION) run(sync_and_run(build_queue(config, db, source, drive))) f = drive.uploaded()[0] topic_folder = drive.files[f.parents[0]] assert topic_folder.name == "Misc• Grammar Practice" assert drive.files[topic_folder.parents[0]].name == "English ( SSC MAINS )" # -- duplicates ---------------------------------------------------------------------------- def test_duplicate_message_detection(config, db, source, drive): msg = source.add(1, EXAMPLE_CAPTION, media_id=999) forwarded = source.add(2, EXAMPLE_CAPTION, video_bytes(700_000, seed=1), media_id=999) queue = build_queue(config, db, source, drive) assert queue.register(msg) == "new" assert queue.register(msg) == "known" assert queue.register(forwarded) == "skipped" assert "duplicate of message 1" in db.get(key(2)).skip_reason run(sync_and_run(queue)) assert len(drive.uploaded()) == 1 def test_filename_collision_gets_suffix(config, db, source, drive): for i in (1, 2, 3): source.add(i, caption("Class 01 Aman sir.mp4", "Misc• Grammar Practice")) run(sync_and_run(build_queue(config, db, source, drive))) assert drive_names(drive) == ["Class 01 Aman sir (2).mp4", "Class 01 Aman sir (3).mp4", "Class 01 Aman sir.mp4"] def test_existing_different_file_is_never_overwritten(config, db, source, drive): queue = build_queue(config, db, source, drive) folder_id = queue.folders.resolve(["Topic"]) existing = drive.add_file(folder_id, "Class 01.mp4", b"someone else's video") source.add(1, caption("Class 01.mp4", "Topic")) run(sync_and_run(queue)) assert drive.files[existing.id].md5 == hashlib.md5(b"someone else's video").hexdigest() assert drive_names(drive) == ["Class 01 (2).mp4", "Class 01.mp4"] @pytest.mark.parametrize( "policy,expected_state,expected_names", [ ("suffix_message_id", State.COMPLETED, ["X [2].mp4", "X.mp4"]), ("skip", State.SKIPPED, ["X.mp4"]), ("fail", State.FAILED, ["X.mp4"]), ], ) def test_duplicate_policies(tmp_path, source, drive, policy, expected_state, expected_names): config = make_config(tmp_path, DUPLICATE_POLICY=policy) db = Database.from_config(config) source.add(1, caption("X.mp4", "T")) source.add(2, caption("X.mp4", "T")) run(sync_and_run(build_queue(config, db, source, drive))) assert db.get(key(2)).state == expected_state assert drive_names(drive) == expected_names def test_identical_existing_file_is_reused_not_duplicated(config, db, source, drive): content = video_bytes(700_000, seed=5) queue = build_queue(config, db, source, drive) folder_id = queue.folders.resolve(["T"]) existing = drive.add_file(folder_id, "X.mp4", content) source.add(1, caption("X.mp4", "T"), content) run(sync_and_run(queue)) assert db.get(key(1)).drive_file_id == existing.id assert len(drive.uploaded()) == 1 and drive.upload_calls == 0 # -- failures and recovery ----------------------------------------------------------------- def test_failed_download_is_retried_and_resumed(config, db, source, drive): content = video_bytes(3 * 1024 * 1024 + 123, seed=9) source.add(1, caption("Big.mp4", "T"), content) source.failures[1] = [("after", 2 * 1024 * 1024, ConnectionError("network dropped"))] queue = build_queue(config, db, source, drive) run(sync_and_run(queue)) rec = db.get(key(1)) assert rec.state == State.COMPLETED and rec.attempts == 1 offsets = [o for m, o in source.download_offsets if m == 1] assert offsets[0] == 0 and offsets[1] == 2 * 1024 * 1024 # resumed, not restarted assert drive.uploaded()[0].md5 == hashlib.md5(content).hexdigest() assert State.RETRY_WAIT.value in [h[1] for h in db.history(key(1))] def test_failed_upload_is_retried(config, db, source, drive): source.add(1, caption("A.mp4", "T")) drive.upload_failures = [TransientError("Drive 503")] run(sync_and_run(build_queue(config, db, source, drive))) assert db.get(key(1)).state == State.COMPLETED assert len(drive.uploaded()) == 1 def test_retries_are_bounded(config, db, source, drive): source.add(1, caption("A.mp4", "T")) source.add(2, caption("B.mp4", "T")) drive.upload_failures = [TransientError("Drive 503")] * config.max_attempts run(sync_and_run(build_queue(config, db, source, drive))) rec = db.get(key(1)) assert rec.state == State.FAILED and rec.attempts == config.max_attempts assert "Drive 503" in rec.last_error assert db.get(key(2)).state == State.COMPLETED # the queue was not blocked def test_permanent_error_fails_without_blocking_queue(config, db, source, drive): source.add(1, caption("Gone.mp4", "T")) source.add(2, caption("B.mp4", "T")) source.failures[1] = [PermanentError("message deleted")] run(sync_and_run(build_queue(config, db, source, drive))) assert db.get(key(1)).state == State.FAILED and db.get(key(1)).attempts == 0 assert db.get(key(2)).state == State.COMPLETED def test_strict_stop_halts_after_failure(tmp_path, source, drive): config = make_config(tmp_path, STRICT_STOP="true") db = Database.from_config(config) source.add(1, caption("Gone.mp4", "T")) source.add(2, caption("B.mp4", "T")) source.failures[1] = [PermanentError("message deleted")] queue = build_queue(config, db, source, drive) run(sync_and_run(queue)) assert queue.status == "halted" assert db.get(key(2)).state == State.PENDING def test_retry_command_requeues_failed(config, db, source, drive): source.add(1, caption("A.mp4", "T")) source.failures[1] = [PermanentError("temporarily missing")] queue = build_queue(config, db, source, drive) run(sync_and_run(queue)) assert db.get(key(1)).state == State.FAILED async def retry(): assert queue.retry_failed() == 1 queue.start() await queue.wait() run(retry()) assert db.get(key(1)).state == State.COMPLETED def test_disk_full_halts_queue_without_using_attempts(config, db, source, drive): source.add(1, caption("A.mp4", "T")) queue = build_queue(config, db, source, drive, disk_usage=lambda p: (100, 100, 10)) run(sync_and_run(queue)) rec = db.get(key(1)) assert queue.status == "halted" and "disk space" in queue.halt_reason assert rec.state == State.RETRY_WAIT and rec.attempts == 0 def test_restart_during_upload_resumes_session(config, source, drive): source.add(1, caption("Big.mp4", "T"), video_bytes(2_000_000, seed=3)) db = Database.from_config(config) queue = build_queue(config, db, source, drive) drive.crash_at_bytes = 700_000 async def crash_run(): await queue.sync() with pytest.raises(SimulatedCrash): await queue.process_item(db.next_item()) run(crash_run()) rec = db.get(key(1)) assert rec.state == State.UPLOADING and rec.upload_session_uri db.close() # "new process" db2 = Database.from_config(config) db2.recover_stale() queue2 = build_queue(config, db2, source, drive) async def resume(): queue2.start() await queue2.wait() run(resume()) assert db2.get(key(1)).state == State.COMPLETED assert drive.resumed and drive.resumed[0][1] >= 700_000 assert len(drive.uploaded()) == 1 assert len([1 for m, _ in source.download_offsets if m == 1]) == 1 # local file reused, no re-download def test_restart_with_wiped_disk_redownloads_and_resumes_upload(config, source, drive): """Heroku dyno restart: ephemeral disk is gone but the DB and upload session survive.""" source.add(1, caption("Big.mp4", "T"), video_bytes(2_000_000, seed=4)) db = Database.from_config(config) queue = build_queue(config, db, source, drive) drive.crash_at_bytes = 1_000_000 async def crash_run(): await queue.sync() with pytest.raises(SimulatedCrash): await queue.process_item(db.next_item()) run(crash_run()) shutil.rmtree(config.download_dir) db2 = Database.from_config(config) db2.recover_stale() queue2 = build_queue(config, db2, source, drive) run(_start_and_wait(queue2)) assert db2.get(key(1)).state == State.COMPLETED assert len([1 for m, _ in source.download_offsets if m == 1]) == 2 assert drive.resumed and drive.resumed[0][1] >= 1_000_000 assert len(drive.uploaded()) == 1 async def _start_and_wait(queue): queue.start() await queue.wait() def test_upload_success_then_database_write_failure(config, source, drive): source.add(1, caption("A.mp4", "T")) db = Database.from_config(config) real_transition = db.transition failed = {"done": False} def flaky_transition(k, state, **kw): if state == State.COMPLETED and not failed["done"]: failed["done"] = True raise db.errors[0]("database is locked") return real_transition(k, state, **kw) db.transition = flaky_transition queue = build_queue(config, db, source, drive) async def first(): await queue.sync() await queue.process_item(db.next_item()) run(first()) rec = db.get(key(1)) assert rec.state == State.UPLOADING # completion was not recorded assert len(drive.uploaded()) == 1 assert (config.download_dir / f"{abs(CHANNEL_ID)}_1.video").exists() # not cleaned up # Same process continues (or a restart): it reconciles instead of re-uploading. run(_start_and_wait(queue)) rec = db.get(key(1)) assert rec.state == State.COMPLETED assert len(drive.uploaded()) == 1 and drive.upload_calls == 1 assert rec.drive_file_id == drive.uploaded()[0].id def test_lost_upload_response_does_not_duplicate(config, db, source, drive): source.add(1, caption("A.mp4", "T")) drive.lose_response_once = True run(sync_and_run(build_queue(config, db, source, drive))) assert db.get(key(1)).state == State.COMPLETED assert len(drive.uploaded()) == 1 # -- control / concurrency ------------------------------------------------------------------ def test_sequential_processing_never_overlaps(config, db, source, drive): for i in range(1, 6): source.add(i, caption(f"Class {i}.mp4", "T")) events = [] source.events = events drive.events_sink = events queue = build_queue(config, db, source, drive) async def go(): await queue.sync() assert queue.start() is True assert queue.start() is False # a second worker is never created await queue.wait() run(go()) assert drive.max_active_uploads == 1 and queue.max_concurrent_uploads == 1 order = [(kind, str(ident).split(":")[-1]) for kind, ident in events] expected = [] for i in range(1, 6): expected += [("download", str(i)), ("upload", str(i))] assert order == expected def test_pause_checkpoint_and_resume(config, db, source, drive): source.add(1, caption("Big.mp4", "T"), video_bytes(3_000_000, seed=7)) queue = build_queue(config, db, source, drive) async def go(): await queue.sync() queue.start() while not (config.download_dir / f"{abs(CHANNEL_ID)}_1.part").exists(): await asyncio.sleep(0) assert queue.pause() for _ in range(200): await asyncio.sleep(0.005) if db.get(key(1)).state == State.PENDING: break assert db.get(key(1)).state == State.PENDING assert queue.status == "paused" assert queue.desired_state() == "paused" assert queue.resume() while db.get(key(1)).state != State.COMPLETED: await asyncio.sleep(0.01) await queue.stop() assert queue.desired_state() == "stopped" run(go()) assert len(drive.uploaded()) == 1 def test_graceful_shutdown_keeps_desired_state_and_recovers(config, source, drive): source.add(1, caption("Big.mp4", "T"), video_bytes(3_000_000, seed=8)) db = Database.from_config(config) queue = build_queue(config, db, source, drive) async def go(): await queue.sync() queue.start() while db.get(key(1)).state != State.UPLOADING: await asyncio.sleep(0.001) await queue.shutdown(timeout=10) run(go()) rec = db.get(key(1)) assert rec.state == State.UPLOADING and rec.upload_session_uri assert queue.desired_state() == "running" # auto-resumes after restart db.close() db2 = Database.from_config(config) db2.recover_stale() queue2 = build_queue(config, db2, source, drive) run(_start_and_wait(queue2)) assert db2.get(key(1)).state == State.COMPLETED assert len(drive.uploaded()) == 1 and drive.resumed def test_progress_messages_are_edited_not_spammed(config, db, source, drive): for i in (1, 2): source.add(i, caption(f"Class {i}.mp4", "Misc• Grammar Practice")) notifier = RecordingNotifier() reporter = ProgressReporter(notifier, [111], interval=1) queue = build_queue(config, db, source, drive, reporter=reporter) run(sync_and_run(queue)) texts = [t for _, t in notifier.sent] progress_messages = [t for t in texts if t.startswith("Processing video") or t == "Queue idle"] assert len(progress_messages) <= 1 # one progress message, then edits assert texts[-1].startswith("Run finished") and "Uploaded: 2" in texts[-1]