drive/tests/test_queue_manager.py

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]