184 lines
6.6 KiB
Python
184 lines
6.6 KiB
Python
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()
|
|
|
|
|
|
def test_upgrade_from_schema_v1_keeps_data(config, db_url):
|
|
"""An existing v1 deployment (before /upload tasks) migrates in place without data loss."""
|
|
import database as dbmod
|
|
|
|
v2 = dbmod.MIGRATIONS[1]
|
|
dbmod.MIGRATIONS.pop() # simulate the old code base
|
|
try:
|
|
old = Database.from_config(config)
|
|
with old._tx() as conn: # a row exactly as the v1 code wrote it
|
|
conn.execute(
|
|
"INSERT INTO videos (channel_id, message_id, state, indexes, title, discovered_at, updated_at) "
|
|
"VALUES (?,?,?,?,?,?,?)",
|
|
(CHANNEL_ID, 5, "DOWNLOADING", '["001", "002"]', "Class 01 Aman sir.mp4", "2026-01-01T00:00:00+00:00",
|
|
"2026-01-01T00:00:00+00:00"),
|
|
)
|
|
assert old.get_kv("schema_version") == "1"
|
|
old.close()
|
|
finally:
|
|
dbmod.MIGRATIONS.append(v2)
|
|
new = Database.from_config(config)
|
|
assert new.get_kv("schema_version") == "2"
|
|
rec = new.get((CHANNEL_ID, 5))
|
|
assert rec.state == State.DOWNLOADING and rec.task_id is None and rec.indexes == ["001", "002"]
|
|
task = new.create_task(111)
|
|
assert task.id == 1 and new.next_item(task_id=task.id) is None
|
|
new.close()
|