drive/tests/test_database.py

from datetime import timedelta

import pytest

from caption_parser import parse_caption
from conftest import CHANNEL_ID, EXAMPLE_CAPTION, FakeSource
from database import Database, InstanceLock, InstanceLockError, InvalidTransition, utcnow
from models import State


def _insert(db, message_id, caption=EXAMPLE_CAPTION):
    msg = FakeSource().add(message_id, caption)
    return msg, db.insert_video(msg, parse_caption(caption))


def test_insert_and_metadata_roundtrip(db):
    msg, inserted = _insert(db, 5)
    assert inserted
    rec = db.get((CHANNEL_ID, 5))
    assert rec.state == State.PENDING
    assert rec.indexes == ["001", "002"]
    assert rec.title == "Class 01 Aman sir.mp4"
    assert rec.topic == "Misc• Grammar Practice"
    assert rec.batch == "English ( SSC MAINS )"
    assert rec.extracted_by == "@PUTINxINDIA"
    assert rec.caption == EXAMPLE_CAPTION
    assert rec.expected_size == msg.media.size
    assert rec.download_status == "pending"


def test_duplicate_message_is_not_inserted_twice(db):
    _, first = _insert(db, 7)
    _, second = _insert(db, 7)
    assert first and not second
    assert db.counts()["PENDING"] == 1


def test_valid_state_transitions_and_history(db):
    _insert(db, 1)
    key = (CHANNEL_ID, 1)
    for state in (State.DOWNLOADING, State.DOWNLOADED, State.UPLOADING, State.COMPLETED):
        rec = db.transition(key, state)
        assert rec.state == state
    assert rec.download_status == "done" and rec.upload_status == "done"
    assert [h[1] for h in db.history(key)] == ["PENDING", "DOWNLOADING", "DOWNLOADED", "UPLOADING", "COMPLETED"]


@pytest.mark.parametrize(
    "path",
    [
        [State.COMPLETED, State.PENDING],  # completed is terminal
        [State.UPLOADING],  # cannot upload before downloading
        [State.DOWNLOADING, State.COMPLETED],  # cannot complete without upload/reconcile
    ],
)
def test_invalid_transitions_rejected(db, path):
    _insert(db, 1)
    key = (CHANNEL_ID, 1)
    with pytest.raises(InvalidTransition):
        for state in path:
            db.transition(key, state)


def test_failure_status_fields(db):
    _insert(db, 1)
    key = (CHANNEL_ID, 1)
    db.transition(key, State.DOWNLOADING)
    rec = db.transition(key, State.RETRY_WAIT, attempts=1, last_error="boom")
    assert rec.download_status == "retry_wait" and rec.attempts == 1 and rec.last_error == "boom"
    db.transition(key, State.DOWNLOADED)
    db.transition(key, State.UPLOADING)
    rec = db.transition(key, State.FAILED)
    assert rec.upload_status == "failed"


def test_unknown_field_rejected(db):
    _insert(db, 1)
    with pytest.raises(ValueError):
        db.update_fields((CHANNEL_ID, 1), message_id=3)


def test_recover_stale(db):
    for i in (1, 2, 3):
        _insert(db, i)
    db.transition((CHANNEL_ID, 1), State.DOWNLOADING)
    db.transition((CHANNEL_ID, 2), State.DOWNLOADING)
    db.transition((CHANNEL_ID, 2), State.DOWNLOADED)
    db.transition((CHANNEL_ID, 2), State.UPLOADING)
    recovered = db.recover_stale()
    assert recovered == {"DOWNLOADING": 1, "UPLOADING": 1}
    assert db.get((CHANNEL_ID, 1)).state == State.PENDING
    assert db.get((CHANNEL_ID, 2)).state == State.UPLOADING  # reconciled by the queue


def test_next_item_prefers_unfinished_then_oldest(db):
    for i in (10, 20, 30):
        _insert(db, i)
    assert db.next_item().message_id == 10
    db.transition((CHANNEL_ID, 30), State.DOWNLOADING)
    db.transition((CHANNEL_ID, 30), State.DOWNLOADED)
    assert db.next_item().message_id == 30
    db.transition((CHANNEL_ID, 30), State.RETRY_WAIT, next_retry_at=(utcnow() + timedelta(hours=1)).isoformat(timespec="seconds"))
    assert db.next_item().message_id == 10
    assert db.next_item(utcnow() + timedelta(hours=2)).message_id == 10
    assert db.earliest_retry_at() is not None


def test_folders_kv_sequence(db):
    assert db.get_folder("root", "misc") is None
    db.put_folder("root", "misc", "Misc", "F1")
    db.put_folder("root", "misc", "Misc", "F2")
    assert db.get_folder("root", "misc") == "F2"
    assert db.folder_count() == 1
    db.raise_int("cursor", 5)
    db.raise_int("cursor", 3)
    assert db.get_int("cursor") == 5
    assert [db.next_sequence() for _ in range(3)] == [1, 2, 3]


def test_state_survives_reopen(config, db):
    _insert(db, 1)
    db.set_kv("queue_desired_state", "running")
    db.close()
    again = Database.from_config(config)
    assert again.get((CHANNEL_ID, 1)).title == "Class 01 Aman sir.mp4"
    assert again.get_kv("queue_desired_state") == "running"
    again.close()


def test_instance_lock_blocks_second_worker(config, db):
    first = InstanceLock(config)
    first.acquire()
    second = InstanceLock(config, wait_seconds=0)
    try:
        if config.database_url:
            with pytest.raises(InstanceLockError):
                second.acquire()
        else:
            # flock is per open file description, so test from a second process
            import subprocess
            import sys

            code = (
                "import fcntl,sys; f=open(sys.argv[1],'w');\n"
                "try:\n fcntl.flock(f, fcntl.LOCK_EX|fcntl.LOCK_NB); print('acquired')\n"
                "except BlockingIOError: print('blocked')"
            )
            out = subprocess.run([sys.executable, "-c", code, str(config.database_path) + ".lock"], capture_output=True, text=True)
            assert out.stdout.strip() == "blocked"
    finally:
        first.release()
        second.release()
    third = InstanceLock(config)
    third.acquire()
    third.release()