448 lines
18 KiB
Python
448 lines
18 KiB
Python
"""Acceptance tests for the interactive /upload workflow (mocked Telegram and Drive)."""
|
||
import asyncio
|
||
from datetime import timedelta
|
||
|
||
import pytest
|
||
|
||
from bot_handlers import CommandRouter
|
||
from conftest import (
|
||
CHANNEL_ID,
|
||
FakeDrive,
|
||
FakeSource,
|
||
RecordingNotifier,
|
||
SimulatedCrash,
|
||
build_queue,
|
||
make_config,
|
||
run,
|
||
video_bytes,
|
||
)
|
||
from database import Database, iso, utcnow
|
||
from errors import PermanentError
|
||
from link_parser import internal_channel_id
|
||
from models import State, TaskStatus
|
||
from progress import ProgressReporter
|
||
|
||
A, B, STRANGER = 111, 222, 999
|
||
INTERNAL = internal_channel_id(CHANNEL_ID) # 1234567890
|
||
OTHER_CHANNEL = -1000999999999 # accessible, but not the configured channel
|
||
|
||
|
||
def link(message_id, internal=INTERNAL):
|
||
return f"https://t.me/c/{internal}/{message_id}"
|
||
|
||
|
||
def caption(title, topic):
|
||
return f"Index: 001\nTitle: {title}\nTopic: {topic}\nBatch: SSC\nExtracted By: @x"
|
||
|
||
|
||
def populate(source, size=600_000):
|
||
"""Messages 13..21; 15 is text, 18 is missing; 13 and 21 lie outside the 14-20 range."""
|
||
source.add(13, caption("Outside Before.mp4", "Topic X"), video_bytes(size, 13))
|
||
source.add(14, caption("Class 01 Aman sir.mp4", "Misc• Grammar Practice"), video_bytes(size, 14))
|
||
source.add_text(15, "Announcement: new classes tomorrow")
|
||
source.add(16, caption("Class 01 Aman sir.mp4", "Error Detection"), video_bytes(size, 16))
|
||
source.add(17, caption("Class 02 Aman sir.mp4", "Misc• Grammar Practice"), video_bytes(size, 17))
|
||
source.add(19, caption("Class 03 Aman sir", "Misc• Grammar Practice"), video_bytes(size, 19))
|
||
source.add(20, "", video_bytes(size, 20))
|
||
source.add(21, caption("Outside After.mp4", "Topic X"), video_bytes(size, 21))
|
||
source.accessible[OTHER_CHANNEL] = "Other Channel"
|
||
|
||
|
||
class Env:
|
||
def __init__(self, config, db, source, drive, notifier=None):
|
||
self.config, self.db, self.source, self.drive = config, db, source, drive
|
||
self.notifier = notifier or RecordingNotifier()
|
||
reporter = ProgressReporter(self.notifier, [A], interval=60)
|
||
self.queue = build_queue(config, db, source, drive, reporter=reporter)
|
||
self.router = CommandRouter(config, db, self.queue)
|
||
|
||
async def say(self, user, *texts):
|
||
replies = []
|
||
for text in texts:
|
||
replies.append(await self.router.handle(user, text))
|
||
return replies
|
||
|
||
def setup(self, user):
|
||
return self.db.tasks_with_status(
|
||
TaskStatus.AWAITING_START, TaskStatus.AWAITING_END, TaskStatus.AWAITING_CONFIRMATION, admin_user_id=user
|
||
)
|
||
|
||
def drive_names(self):
|
||
return sorted((self.drive.files[f.parents[0]].name, f.name) for f in self.drive.uploaded())
|
||
|
||
|
||
@pytest.fixture
|
||
def env(tmp_path, db_url):
|
||
extra = {"DATABASE_URL": db_url} if db_url else {}
|
||
config = make_config(tmp_path, ADMIN_USER_IDS=f"{A},{B}", **extra)
|
||
db = Database.from_config(config)
|
||
source = FakeSource()
|
||
populate(source)
|
||
e = Env(config, db, source, FakeDrive())
|
||
yield e
|
||
db.close()
|
||
|
||
|
||
async def full_task(env, start=14, end=20, user=A):
|
||
replies = await env.say(user, "/upload", link(start), link(end), "/confirm")
|
||
await env.queue.wait()
|
||
return replies
|
||
|
||
|
||
# 1-4, 7-10: conversation, inclusive ascending range, metadata, skips ------------------------
|
||
|
||
|
||
def test_upload_conversation_and_range_processing(env):
|
||
async def go():
|
||
r = await env.say(A, "/upload")
|
||
assert "📥 <b>Start Upload Task</b>" in r[0] and "https://t.me/c/3909337424/14" in r[0]
|
||
assert env.setup(A)[0].status == TaskStatus.AWAITING_START
|
||
|
||
r = await env.say(A, link(14))
|
||
assert "📌 <b>Enter the End Link</b>" in r[0]
|
||
assert env.setup(A)[0].start_message_id == 14
|
||
|
||
r = await env.say(A, link(20))
|
||
assert "✅ <b>Upload Task Ready</b>" in r[0]
|
||
assert "Start Message ID: 14" in r[0] and "End Message ID: 20" in r[0]
|
||
assert "Test Channel (configured source channel)" in r[0]
|
||
assert env.source.download_offsets == [] # nothing downloaded before /confirm
|
||
|
||
r = await env.say(A, "/confirm")
|
||
assert "Upload Task #1 confirmed" in r[0]
|
||
await env.queue.wait()
|
||
|
||
run(go())
|
||
task = env.db.get_task(1)
|
||
assert task.status == TaskStatus.COMPLETED
|
||
# inclusive range 14..20, nothing outside it
|
||
assert [m for m, _ in env.db.task_items(1)] == [14, 15, 16, 17, 18, 19, 20]
|
||
downloaded = [m for kind, m in env.source.events if kind == "download"]
|
||
assert downloaded == [14, 16, 17, 19, 20] # ascending, videos only
|
||
assert not env.db.exists(CHANNEL_ID, 13) and not env.db.exists(CHANNEL_ID, 21)
|
||
# caption metadata -> folders and filenames
|
||
assert env.drive_names() == [
|
||
("Error Detection", "Class 01 Aman sir.mp4"),
|
||
("Misc• Grammar Practice", "Class 01 Aman sir.mp4"),
|
||
("Misc• Grammar Practice", "Class 02 Aman sir.mp4"),
|
||
("Misc• Grammar Practice", "Class 03 Aman sir.mp4"),
|
||
("Uncategorized", "telegram_20.mp4"),
|
||
]
|
||
# non-video and missing messages are skipped, not failed
|
||
assert dict(env.db.task_items(1))[15] == "non_video"
|
||
assert dict(env.db.task_items(1))[18] == "missing"
|
||
assert (task.uploaded_count, task.failed_count, task.non_video_count, task.missing_count) == (5, 0, 1, 1)
|
||
assert task.examined_count == 6 and task.eligible_count == 5
|
||
assert task.deleted_files_count == 5
|
||
assert list(env.config.download_dir.iterdir()) == []
|
||
summary = [t for _, t in env.notifier.sent if "Upload Task Completed" in t][0]
|
||
for fragment in ("Requested range: 14–20", "Eligible videos found: 5", "Uploaded: 5", "Duplicates skipped: 0",
|
||
"Non-video messages skipped: 1", "Failed: 0", "Local temporary files deleted: 5",
|
||
"Misc• Grammar Practice", "Error Detection"):
|
||
assert fragment in summary
|
||
|
||
|
||
def test_confirm_requires_both_validated_links(env):
|
||
async def go():
|
||
assert "no upload task waiting" in (await env.say(A, "/confirm"))[0]
|
||
await env.say(A, "/upload")
|
||
assert "Both the start and end links" in (await env.say(A, "/confirm"))[0]
|
||
await env.say(A, link(14))
|
||
assert "Both the start and end links" in (await env.say(A, "/confirm"))[0]
|
||
assert env.db.next_runnable_task() is None
|
||
|
||
run(go())
|
||
assert env.source.download_offsets == []
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"bad,fragment",
|
||
[
|
||
("example.com/post/14", "not a Telegram"),
|
||
("https://t.me/+invite", "Invite links"),
|
||
("https://t.me/c/1234567890", "must look like"),
|
||
("https://t.me/c/555/14", "not accessible"), # channel the account cannot access
|
||
(link(14, internal_channel_id(OTHER_CHANNEL)), "not from the configured source channel"),
|
||
(link(18), "not found"), # deleted message
|
||
],
|
||
)
|
||
def test_invalid_start_links_rejected(env, bad, fragment):
|
||
replies = run(env.say(A, "/upload", bad))
|
||
assert replies[1].startswith("❌") and fragment in replies[1]
|
||
assert env.setup(A)[0].status == TaskStatus.AWAITING_START
|
||
|
||
|
||
def test_end_link_errors_keep_start_link(env):
|
||
replies = run(env.say(
|
||
A, "/upload", link(14),
|
||
link(20, internal_channel_id(OTHER_CHANNEL)), # different channel
|
||
link(13), # before start
|
||
"https://t.me/c/1234567890/abc", # malformed
|
||
))
|
||
assert "different channel" in replies[2]
|
||
assert "greater than or equal" in replies[3]
|
||
assert replies[4].startswith("❌")
|
||
setup = env.setup(A)[0]
|
||
assert setup.status == TaskStatus.AWAITING_END and setup.start_message_id == 14
|
||
assert "Upload Task Ready" in run(env.say(A, link(14)))[0] # single-message range is allowed
|
||
|
||
|
||
def test_other_channels_allowed_when_unrestricted(tmp_path):
|
||
config = make_config(tmp_path, ADMIN_USER_IDS=str(A), UPLOAD_RESTRICT_TO_SOURCE_CHANNEL="false")
|
||
db = Database.from_config(config)
|
||
source = FakeSource()
|
||
populate(source)
|
||
e = Env(config, db, source, FakeDrive())
|
||
|
||
async def other_exists(channel_id, message_id):
|
||
return True
|
||
|
||
source.message_accessible = other_exists
|
||
replies = run(e.say(A, "/upload", link(1, internal_channel_id(OTHER_CHANNEL)), link(5, internal_channel_id(OTHER_CHANNEL))))
|
||
assert "Channel: Other Channel" in replies[2]
|
||
|
||
|
||
def test_range_size_limit(tmp_path):
|
||
config = make_config(tmp_path, ADMIN_USER_IDS=str(A), UPLOAD_MAX_RANGE="3")
|
||
e = Env(config, Database.from_config(config), FakeSource(), FakeDrive())
|
||
populate(e.source)
|
||
replies = run(e.say(A, "/upload", link(14), link(20)))
|
||
assert "maximum is 3" in replies[2]
|
||
|
||
|
||
def test_non_admins_cannot_use_upload_commands(env):
|
||
for command in ("/upload", "/confirm", "/cancel", "/taskstatus", "/uploadhistory"):
|
||
assert "not authorized" in run(env.say(STRANGER, command))[0]
|
||
assert run(env.say(STRANGER, link(14)))[0] is None # plain text from non-admins is ignored
|
||
assert env.db.recent_tasks() == [] and env.setup(STRANGER) == []
|
||
|
||
|
||
# 15: administrators are isolated ------------------------------------------------------
|
||
|
||
|
||
def test_admin_setups_are_isolated(env):
|
||
async def go():
|
||
await env.say(A, "/upload", link(14))
|
||
r = await env.say(B, link(16))
|
||
assert "send /upload first" in r[0]
|
||
await env.say(B, "/upload", link(16))
|
||
assert env.setup(A)[0].start_message_id == 14
|
||
assert env.setup(B)[0].start_message_id == 16
|
||
r = await env.say(B, "/cancel")
|
||
assert "cancelled" in r[0]
|
||
assert env.setup(B) == []
|
||
assert env.setup(A)[0].status == TaskStatus.AWAITING_END # untouched
|
||
await env.say(A, link(17), "/confirm")
|
||
r = await env.say(B, "/cancel")
|
||
assert "belong to another administrator" in r[0]
|
||
await env.queue.wait()
|
||
|
||
run(go())
|
||
assert env.db.get_task(1).status == TaskStatus.COMPLETED
|
||
|
||
|
||
def test_second_upload_does_not_overwrite_setup(env):
|
||
replies = run(env.say(A, "/upload", link(14), "/upload"))
|
||
assert "already have an upload setup" in replies[2]
|
||
setup = env.setup(A)
|
||
assert len(setup) == 1 and setup[0].start_message_id == 14
|
||
|
||
|
||
# setup persistence / expiry ------------------------------------------------------------
|
||
|
||
|
||
def test_setup_survives_restart(env):
|
||
run(env.say(A, "/upload", link(14)))
|
||
# new process: fresh queue/router on the same durable database
|
||
again = Env(env.config, env.db, env.source, env.drive)
|
||
replies = run(again.say(A, link(20)))
|
||
assert "Upload Task Ready" in replies[0]
|
||
|
||
|
||
def test_stale_setup_expires(env):
|
||
run(env.say(A, "/upload", link(14)))
|
||
setup = env.setup(A)[0]
|
||
env.db.update_task(setup.id, updated_at=iso(utcnow() - timedelta(minutes=env.config.upload_setup_timeout + 1)))
|
||
reply = run(env.say(A, "/confirm"))[0]
|
||
assert "expired" in reply
|
||
assert env.db.get_task(setup.id).status == TaskStatus.CANCELLED
|
||
assert "Start Upload Task" in run(env.say(A, "/upload"))[0]
|
||
|
||
|
||
# 11, 12: cleanup ordering and failures -------------------------------------------------
|
||
|
||
|
||
def test_local_file_deleted_only_after_verified_upload_is_recorded(env):
|
||
seen = []
|
||
real_cleanup = env.queue.downloader.cleanup
|
||
|
||
def spy(rec):
|
||
current = env.db.get(rec.key)
|
||
if current.state == State.COMPLETED:
|
||
assert current.drive_file_id in env.drive.files
|
||
assert env.drive.files[current.drive_file_id].size == current.local_size
|
||
assert env.config.download_dir.joinpath(f"{abs(CHANNEL_ID)}_{rec.message_id}.video").exists()
|
||
seen.append((rec.message_id, current.state))
|
||
return real_cleanup(rec)
|
||
|
||
env.queue.downloader.cleanup = spy
|
||
run(full_task(env))
|
||
assert seen and all(state == State.COMPLETED for _, state in seen)
|
||
|
||
|
||
def test_failed_upload_is_not_completed(env):
|
||
env.drive.upload_failures = [PermanentError("insufficientFilePermissions")]
|
||
run(full_task(env))
|
||
rec = env.db.get((CHANNEL_ID, 14))
|
||
assert rec.state == State.FAILED and rec.drive_file_id is None
|
||
task = env.db.get_task(1)
|
||
assert task.status == TaskStatus.COMPLETED and task.failed_count == 1 and task.uploaded_count == 4
|
||
assert dict(env.db.task_items(1))[14] == "failed"
|
||
assert not any(p.name.startswith(f"{abs(CHANNEL_ID)}_14.") for p in env.config.download_dir.iterdir())
|
||
|
||
|
||
def test_failed_download_file_can_be_kept(tmp_path):
|
||
config = make_config(tmp_path, ADMIN_USER_IDS=str(A), KEEP_FAILED_DOWNLOADS="true")
|
||
e = Env(config, Database.from_config(config), FakeSource(), FakeDrive())
|
||
populate(e.source)
|
||
e.drive.upload_failures = [PermanentError("denied")]
|
||
run(full_task(e))
|
||
assert (config.download_dir / f"{abs(CHANNEL_ID)}_14.video").exists()
|
||
|
||
|
||
# 14: repeating a task ------------------------------------------------------------------
|
||
|
||
|
||
def test_repeating_completed_task_does_not_duplicate(env):
|
||
run(full_task(env))
|
||
calls = env.drive.upload_calls
|
||
run(full_task(env))
|
||
second = env.db.get_task(2)
|
||
assert second.status == TaskStatus.COMPLETED
|
||
assert second.uploaded_count == 0 and second.duplicate_count == 5
|
||
assert env.drive.upload_calls == calls and len(env.drive.uploaded()) == 5
|
||
|
||
|
||
# 16: cancellation ----------------------------------------------------------------------
|
||
|
||
|
||
def test_cancel_running_task_keeps_completed_work(tmp_path):
|
||
config = make_config(tmp_path, ADMIN_USER_IDS=str(A))
|
||
db = Database.from_config(config)
|
||
source = FakeSource()
|
||
populate(source, size=3_000_000)
|
||
e = Env(config, db, source, FakeDrive())
|
||
part16 = config.download_dir / f"{abs(CHANNEL_ID)}_16.part"
|
||
|
||
async def go():
|
||
await e.say(A, "/upload", link(14), link(20), "/confirm")
|
||
while not part16.exists():
|
||
await asyncio.sleep(0.001)
|
||
reply = (await e.say(A, "/cancel"))[0]
|
||
assert "cancelled at a safe checkpoint" in reply
|
||
await e.queue.wait()
|
||
|
||
run(go())
|
||
task = db.get_task(1)
|
||
assert task.status == TaskStatus.CANCELLED
|
||
assert db.get((CHANNEL_ID, 14)).state == State.COMPLETED
|
||
assert len(e.drive.uploaded()) == 1 # completed upload preserved in Drive
|
||
for message_id in (16, 17, 19, 20):
|
||
rec = db.get((CHANNEL_ID, message_id))
|
||
assert rec.state == State.SKIPPED and rec.skip_reason.startswith("cancelled:")
|
||
assert list(config.download_dir.iterdir()) == []
|
||
assert "Upload Task Cancelled" in "".join(t for _, t in e.notifier.sent)
|
||
|
||
# the same range can be requested again: cancelled items are re-queued, done ones are not re-uploaded
|
||
run(full_task(e))
|
||
second = db.get_task(2)
|
||
assert second.uploaded_count == 4 and second.duplicate_count == 1
|
||
assert len(e.drive.uploaded()) == 5
|
||
|
||
|
||
def test_cancel_queued_task_and_explicit_id(env):
|
||
async def go():
|
||
await env.say(A, "/pause") # nothing running -> nothing to pause
|
||
env.queue._persist_desired("paused")
|
||
await env.say(A, "/upload", link(14), link(20), "/confirm")
|
||
assert env.db.get_task(1).status == TaskStatus.QUEUED
|
||
reply = (await env.say(B, "/cancel 1"))[0]
|
||
assert "Task #1 cancelled" in reply
|
||
assert "already finished" in (await env.say(A, "/cancel 1"))[0]
|
||
|
||
run(go())
|
||
assert env.db.get_task(1).status == TaskStatus.CANCELLED
|
||
assert env.source.download_offsets == []
|
||
|
||
|
||
# 13, 18: restart recovery ----------------------------------------------------------------
|
||
|
||
|
||
def test_task_recovers_after_crash_and_restart(tmp_path, db_url):
|
||
extra = {"DATABASE_URL": db_url} if db_url else {}
|
||
config = make_config(tmp_path, ADMIN_USER_IDS=str(A), **extra)
|
||
source = FakeSource()
|
||
populate(source, size=1_500_000)
|
||
drive = FakeDrive()
|
||
db = Database.from_config(config)
|
||
e = Env(config, db, source, drive)
|
||
drive.crash_at_bytes = 700_000 # the process is killed half-way through the first upload
|
||
|
||
async def first_run():
|
||
await e.say(A, "/upload", link(14), link(20), "/confirm")
|
||
with pytest.raises(SimulatedCrash):
|
||
await e.queue.wait()
|
||
|
||
run(first_run())
|
||
assert db.get_task(1).status == TaskStatus.RUNNING
|
||
assert db.get((CHANNEL_ID, 14)).state == State.UPLOADING
|
||
db.close()
|
||
|
||
# Heroku restarts the worker: new process, same durable database
|
||
db2 = Database.from_config(config)
|
||
db2.recover_stale()
|
||
e2 = Env(config, db2, source, drive)
|
||
|
||
async def second_run():
|
||
assert e2.queue.boot() == "running"
|
||
await e2.queue.wait()
|
||
|
||
run(second_run())
|
||
task = db2.get_task(1)
|
||
assert task.status == TaskStatus.COMPLETED and task.uploaded_count == 5
|
||
assert len(drive.uploaded()) == 5 # no duplicates
|
||
assert drive.resumed # the interrupted upload session was resumed
|
||
|
||
|
||
def test_pause_resume_task_and_paused_boot(env):
|
||
async def go():
|
||
await env.say(A, "/upload", link(14), link(20), "/confirm")
|
||
while env.db.get((CHANNEL_ID, 14)) is None or env.db.get((CHANNEL_ID, 14)).state != State.COMPLETED:
|
||
await asyncio.sleep(0.001)
|
||
assert "Pausing" in (await env.say(A, "/pause"))[0]
|
||
await asyncio.sleep(0.05)
|
||
assert env.db.get_task(1).status == TaskStatus.PAUSED
|
||
# a restart while paused stays paused
|
||
assert Env(env.config, env.db, env.source, env.drive).queue.boot() == "paused"
|
||
assert "resumed" in (await env.say(A, "/resume"))[0]
|
||
await env.queue.wait()
|
||
|
||
run(go())
|
||
assert env.db.get_task(1).status == TaskStatus.COMPLETED
|
||
|
||
|
||
def test_status_and_history_commands(env):
|
||
async def go():
|
||
await env.say(A, "/upload", link(14))
|
||
status = (await env.say(A, "/taskstatus"))[0]
|
||
assert "waiting for the end link" in status
|
||
await env.say(A, link(20), "/confirm")
|
||
await env.queue.wait()
|
||
status = (await env.say(A, "/taskstatus 1"))[0]
|
||
assert "Task #1</b> — COMPLETED" in status and "Range: 14–20" in status
|
||
history = (await env.say(A, "/uploadhistory"))[0]
|
||
assert "#1 14–20 COMPLETED — uploaded 5" in history
|
||
assert "Upload tasks" not in (await env.say(A, "/status"))[0]
|
||
|
||
run(go())
|