first commit
This commit is contained in:
gosangam
@@ -0,0 +1,373 @@
|
||||
"""Shared fixtures and in-memory fakes for Telegram and Google Drive.
|
||||
|
||||
No test needs real credentials or network access. Database-backed tests run against
|
||||
SQLite and, when the optional ``pgserver`` dev package is installed, a throwaway
|
||||
local PostgreSQL server as well.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import itertools
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from config import Config
|
||||
from database import Database
|
||||
from downloader import Downloader
|
||||
from drive_client import FOLDER_MIME, SOURCE_PROPERTY
|
||||
from errors import Interrupted
|
||||
from models import DriveFile, MediaInfo, SourceMessage
|
||||
|
||||
CHANNEL_ID = -1001234567890
|
||||
EXAMPLE_CAPTION = """Index: 001
|
||||
|
||||
Index: 002
|
||||
|
||||
Title: Class 01 Aman sir.mp4
|
||||
|
||||
Topic: Misc• Grammar Practice
|
||||
|
||||
Batch: English ( SSC MAINS )
|
||||
|
||||
Extracted By: @PUTINxINDIA"""
|
||||
|
||||
|
||||
def make_config(tmp_path: Path, **env) -> Config:
|
||||
base = {
|
||||
"TELEGRAM_API_ID": "12345",
|
||||
"TELEGRAM_API_HASH": "0123456789abcdef0123456789abcdef",
|
||||
"TELEGRAM_BOT_TOKEN": "123456:TEST-TOKEN-VALUE",
|
||||
"ADMIN_USER_IDS": "111",
|
||||
"SOURCE_CHANNEL_ID": str(CHANNEL_ID),
|
||||
"DATABASE_PATH": str(tmp_path / "bot.db"),
|
||||
"DOWNLOAD_DIR": str(tmp_path / "downloads"),
|
||||
"SESSION_DIR": str(tmp_path / "sessions"),
|
||||
"LOG_DIR": str(tmp_path / "logs"),
|
||||
"MIN_FREE_DISK_SPACE": "0",
|
||||
"RETRY_DELAY": "0",
|
||||
"PROCESSING_DELAY": "0",
|
||||
"MAX_RETRIES": "3",
|
||||
"UPLOAD_CHUNK_SIZE": "256KB",
|
||||
}
|
||||
base.update({k: str(v) for k, v in env.items()})
|
||||
return Config.from_env(base)
|
||||
|
||||
|
||||
def video_bytes(size: int, seed: int = 0) -> bytes:
|
||||
"""Bytes that look like an MP4 container (ftyp box) followed by filler."""
|
||||
header = b"\x00\x00\x00\x18ftypisom\x00\x00\x02\x00isomiso2"
|
||||
filler = hashlib.sha256(str(seed).encode()).digest()
|
||||
body = (filler * (size // len(filler) + 1))[: size - len(header)]
|
||||
return header + body
|
||||
|
||||
|
||||
# -- PostgreSQL (optional) ---------------------------------------------------------------
|
||||
|
||||
_PG_SERVER = None
|
||||
|
||||
|
||||
def _pg_uri() -> str | None:
|
||||
global _PG_SERVER
|
||||
try:
|
||||
import pgserver
|
||||
except ImportError:
|
||||
return None
|
||||
if _PG_SERVER is None:
|
||||
_PG_SERVER = pgserver.get_server(tempfile.mkdtemp(prefix="tdb-pg-"), cleanup_mode="stop")
|
||||
return _PG_SERVER.get_uri()
|
||||
|
||||
|
||||
_pg_counter = itertools.count()
|
||||
|
||||
|
||||
def _fresh_pg_url() -> str:
|
||||
import psycopg
|
||||
|
||||
uri = _pg_uri()
|
||||
name = f"test_{os.getpid()}_{next(_pg_counter)}"
|
||||
with psycopg.connect(uri, autocommit=True) as conn:
|
||||
conn.execute(f"CREATE DATABASE {name}")
|
||||
return uri.replace("/postgres?", f"/{name}?")
|
||||
|
||||
|
||||
BACKENDS = ["sqlite"] + (["postgres"] if _pg_uri() else [])
|
||||
|
||||
|
||||
@pytest.fixture(params=BACKENDS)
|
||||
def backend(request):
|
||||
return request.param
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db_url(backend):
|
||||
return _fresh_pg_url() if backend == "postgres" else None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config(tmp_path, db_url):
|
||||
env = {"DATABASE_URL": db_url} if db_url else {}
|
||||
return make_config(tmp_path, **env)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db(config):
|
||||
database = Database.from_config(config)
|
||||
yield database
|
||||
database.close()
|
||||
|
||||
|
||||
# -- fakes -------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FakeSource:
|
||||
CHUNK = 512 * 1024
|
||||
|
||||
def __init__(self, channel_id: int = CHANNEL_ID):
|
||||
self.channel_id = channel_id
|
||||
self.channel_title = "Test Channel"
|
||||
# Other channels the fake account can access: marked id -> title, plus usernames.
|
||||
self.accessible = {channel_id: self.channel_title}
|
||||
self.usernames = {"test_channel": channel_id}
|
||||
self.messages: dict[int, tuple[SourceMessage, bytes]] = {}
|
||||
self.failures: dict[int, list] = {}
|
||||
self.download_offsets: list[tuple[int, int]] = []
|
||||
self.events: list[tuple[str, int]] = []
|
||||
|
||||
def add(self, message_id, caption, content: bytes | None = None, *, media_id=None, kind="video", mime="video/mp4", file_name=None):
|
||||
content = content if content is not None else video_bytes(700_000, seed=message_id)
|
||||
media = MediaInfo(kind, len(content), mime, file_name, media_id if media_id is not None else 10_000 + message_id)
|
||||
msg = SourceMessage(self.channel_id, message_id, caption, media)
|
||||
self.messages[message_id] = (msg, content)
|
||||
return msg
|
||||
|
||||
def add_text(self, message_id, text):
|
||||
msg = SourceMessage(self.channel_id, message_id, text, None)
|
||||
self.messages[message_id] = (msg, b"")
|
||||
return msg
|
||||
|
||||
async def iter_history(self, *, min_id, start_date):
|
||||
for message_id in sorted(self.messages):
|
||||
if message_id >= min_id:
|
||||
yield self.messages[message_id][0]
|
||||
|
||||
async def resolve_link(self, link):
|
||||
from errors import PermanentError
|
||||
|
||||
if link.kind == "private":
|
||||
channel_id = int(link.channel_ref)
|
||||
else:
|
||||
channel_id = self.usernames.get(link.channel.lower())
|
||||
if channel_id not in self.accessible:
|
||||
raise PermanentError("no access", public="This channel is not accessible to the bot's Telegram account.")
|
||||
return channel_id, self.accessible[channel_id]
|
||||
|
||||
async def message_accessible(self, channel_id, message_id):
|
||||
return channel_id == self.channel_id and message_id in self.messages
|
||||
|
||||
async def iter_range(self, channel_id, start_id, end_id):
|
||||
self.events.append(("scan", (start_id, end_id)))
|
||||
for message_id in range(start_id, end_id + 1):
|
||||
entry = self.messages.get(message_id) if channel_id == self.channel_id else None
|
||||
yield message_id, (entry[0] if entry else None)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def iter_media_chunks(self, message_id, *, offset, expected_size, channel_id=None):
|
||||
assert channel_id in (None, self.channel_id)
|
||||
self.download_offsets.append((message_id, offset))
|
||||
self.events.append(("download", message_id))
|
||||
content = self.messages[message_id][1]
|
||||
plan = self.failures.get(message_id) or []
|
||||
failure = plan.pop(0) if plan else None
|
||||
sent = 0
|
||||
for start in range(offset, len(content), self.CHUNK):
|
||||
if isinstance(failure, tuple) and failure[0] == "after" and sent >= failure[1]:
|
||||
raise failure[2]
|
||||
if isinstance(failure, BaseException):
|
||||
raise failure
|
||||
chunk = content[start : start + self.CHUNK]
|
||||
sent += len(chunk)
|
||||
yield chunk
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
class SimulatedCrash(BaseException):
|
||||
"""Stands in for the process being killed (not caught by `except Exception`)."""
|
||||
|
||||
|
||||
class FakeDrive:
|
||||
CHUNK = 256 * 1024
|
||||
|
||||
def __init__(self):
|
||||
self.files: dict[str, DriveFile] = {}
|
||||
self.sessions: dict[str, dict] = {}
|
||||
self._ids = itertools.count(1)
|
||||
self._lock = threading.Lock()
|
||||
self.active_uploads = 0
|
||||
self.max_active_uploads = 0
|
||||
self.upload_failures: list[BaseException] = []
|
||||
self.crash_at_bytes: int | None = None
|
||||
self.lose_response_once = False
|
||||
self.resumed: list[tuple[str, int]] = []
|
||||
self.upload_calls = 0
|
||||
self.events_sink: list | None = None
|
||||
self.created_folders = 0
|
||||
|
||||
def _id(self, prefix):
|
||||
return f"{prefix}{next(self._ids)}"
|
||||
|
||||
def add_folder(self, parent, name):
|
||||
folder = DriveFile(self._id("fold"), name, parents=(parent,), mime_type=FOLDER_MIME)
|
||||
self.files[folder.id] = folder
|
||||
return folder
|
||||
|
||||
def add_file(self, parent, name, content: bytes, app_properties=None):
|
||||
f = DriveFile(
|
||||
self._id("file"), name, size=len(content), md5=hashlib.md5(content).hexdigest(), parents=(parent,),
|
||||
mime_type="video/mp4", app_properties=dict(app_properties or {}),
|
||||
)
|
||||
self.files[f.id] = f
|
||||
return f
|
||||
|
||||
def uploaded(self):
|
||||
return [f for f in self.files.values() if f.mime_type != FOLDER_MIME and not f.trashed]
|
||||
|
||||
# DriveClient interface
|
||||
def get_file(self, file_id):
|
||||
return self.files.get(file_id)
|
||||
|
||||
def find_folders(self, parent_id, name):
|
||||
return [f for f in self.files.values() if f.mime_type == FOLDER_MIME and parent_id in f.parents and f.name == name and not f.trashed]
|
||||
|
||||
def create_folder(self, parent_id, name):
|
||||
self.created_folders += 1
|
||||
return self.add_folder(parent_id, name)
|
||||
|
||||
def find_by_name(self, parent_id, name):
|
||||
return [f for f in self.uploaded() if parent_id in f.parents and f.name == name]
|
||||
|
||||
def find_by_source_key(self, key):
|
||||
return [f for f in self.uploaded() if f.app_properties.get(SOURCE_PROPERTY) == key]
|
||||
|
||||
def share_with_link(self, file_id):
|
||||
pass
|
||||
|
||||
def upload_file(self, path, *, name, parent_id, mime_type, size, app_properties, description=None,
|
||||
session_uri, on_session, on_progress, should_stop):
|
||||
with self._lock:
|
||||
self.active_uploads += 1
|
||||
self.max_active_uploads = max(self.max_active_uploads, self.active_uploads)
|
||||
self.upload_calls += 1
|
||||
if self.events_sink is not None:
|
||||
self.events_sink.append(("upload", app_properties[SOURCE_PROPERTY]))
|
||||
try:
|
||||
data = Path(path).read_bytes()
|
||||
assert len(data) == size
|
||||
session = self.sessions.get(session_uri) if session_uri else None
|
||||
if session is not None:
|
||||
self.resumed.append((session_uri, len(session["data"])))
|
||||
if session.get("file"):
|
||||
return session["file"]
|
||||
else:
|
||||
session_uri = f"https://upload.example/session/{self._id('s')}"
|
||||
session = {"data": b"", "name": name, "parent": parent_id, "props": app_properties}
|
||||
self.sessions[session_uri] = session
|
||||
on_session(session_uri)
|
||||
if self.upload_failures:
|
||||
raise self.upload_failures.pop(0)
|
||||
while len(session["data"]) < size:
|
||||
if should_stop():
|
||||
raise Interrupted("upload paused")
|
||||
start = len(session["data"])
|
||||
session["data"] += data[start : start + self.CHUNK]
|
||||
on_progress(len(session["data"]), size)
|
||||
time.sleep(0.001)
|
||||
if self.crash_at_bytes is not None and len(session["data"]) >= self.crash_at_bytes:
|
||||
self.crash_at_bytes = None
|
||||
raise SimulatedCrash()
|
||||
f = self.add_file(session["parent"], session["name"], session["data"], session["props"])
|
||||
session["file"] = f
|
||||
if self.lose_response_once:
|
||||
self.lose_response_once = False
|
||||
from errors import TransientError
|
||||
|
||||
raise TransientError("connection reset after upload completed")
|
||||
return f
|
||||
finally:
|
||||
with self._lock:
|
||||
self.active_uploads -= 1
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def source():
|
||||
return FakeSource()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def drive():
|
||||
return FakeDrive()
|
||||
|
||||
|
||||
class RecordingNotifier:
|
||||
def __init__(self):
|
||||
self.sent: list[tuple[int, str]] = []
|
||||
self.edits: list[tuple[int, int, str]] = []
|
||||
self._ids = itertools.count(1)
|
||||
|
||||
async def send(self, chat_id, text):
|
||||
self.sent.append((chat_id, text))
|
||||
return next(self._ids)
|
||||
|
||||
async def edit(self, chat_id, message_id, text):
|
||||
self.edits.append((chat_id, message_id, text))
|
||||
|
||||
|
||||
def build_queue(config, db, source, drive, *, reporter=None, disk_usage=None, root="root-folder"):
|
||||
from drive_client import FolderResolver
|
||||
from queue_manager import QueueManager
|
||||
|
||||
kwargs = {"disk_usage": disk_usage} if disk_usage else {}
|
||||
downloader = Downloader(
|
||||
source, config.download_dir, min_free_bytes=config.min_free_disk_space,
|
||||
max_temp_bytes=config.max_temp_storage, compute_sha256=config.compute_sha256, **kwargs,
|
||||
)
|
||||
if root not in drive.files:
|
||||
drive.files[root] = DriveFile(root, "Root", mime_type=FOLDER_MIME)
|
||||
resolver = FolderResolver(drive, db, root, sleep=lambda s: None)
|
||||
return QueueManager(config, db, drive, downloader, resolver, source=source, reporter=reporter, db_error_delay=0)
|
||||
|
||||
|
||||
def run(coro):
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _no_dotenv(monkeypatch):
|
||||
# Never pick up a developer's real .env during tests.
|
||||
monkeypatch.setattr("dotenv.load_dotenv", lambda *a, **k: False)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _restore_logging():
|
||||
import logging
|
||||
|
||||
root = logging.getLogger()
|
||||
handlers, level = list(root.handlers), root.level
|
||||
yield
|
||||
root.handlers[:] = handlers
|
||||
root.setLevel(level)
|
||||
|
||||
|
||||
def pytest_sessionfinish(session, exitstatus):
|
||||
if _PG_SERVER is not None:
|
||||
try:
|
||||
_PG_SERVER.cleanup()
|
||||
except Exception:
|
||||
pass
|
||||
shutil.rmtree(_PG_SERVER.pgdata, ignore_errors=True)
|
||||
@@ -0,0 +1,69 @@
|
||||
import pytest
|
||||
|
||||
from bot_handlers import CommandRouter
|
||||
from conftest import CHANNEL_ID, EXAMPLE_CAPTION, build_queue, run
|
||||
from models import State
|
||||
|
||||
ADMIN, STRANGER = 111, 999
|
||||
|
||||
COMMANDS = ["/start", "/help", "/status", "/run", "/pause", "/resume", "/stop", "/retry", "/failed",
|
||||
"/progress", "/folders", "/check 1", "/sync"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def router(config, db, source, drive):
|
||||
source.add(1, EXAMPLE_CAPTION)
|
||||
queue = build_queue(config, db, source, drive)
|
||||
return CommandRouter(config, db, queue)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("command", COMMANDS)
|
||||
def test_non_admin_is_rejected(router, command):
|
||||
reply = run(router.handle(STRANGER, command))
|
||||
assert reply == "You are not authorized to control this bot."
|
||||
assert router.queue.running is False
|
||||
assert router.db.counts()["PENDING"] == 0 # /sync did not run
|
||||
|
||||
|
||||
def test_none_sender_is_rejected(router):
|
||||
assert "not authorized" in run(router.handle(None, "/status"))
|
||||
|
||||
|
||||
def test_non_commands_are_ignored(router):
|
||||
assert run(router.handle(ADMIN, "hello")) is None
|
||||
|
||||
|
||||
def test_admin_flow(router, drive):
|
||||
async def flow():
|
||||
replies = {}
|
||||
replies["help"] = await router.handle(ADMIN, "/help")
|
||||
replies["sync"] = await router.handle(ADMIN, "/sync@MyBot")
|
||||
replies["check"] = await router.handle(ADMIN, "/check 1")
|
||||
replies["run"] = await router.handle(ADMIN, "/run")
|
||||
await router.queue.wait()
|
||||
replies["progress"] = await router.handle(ADMIN, "/progress")
|
||||
replies["folders"] = await router.handle(ADMIN, "/folders")
|
||||
replies["status"] = await router.handle(ADMIN, "/status")
|
||||
replies["failed"] = await router.handle(ADMIN, "/failed")
|
||||
replies["unknown"] = await router.handle(ADMIN, "/nope")
|
||||
replies["badcheck"] = await router.handle(ADMIN, "/check abc")
|
||||
return replies
|
||||
|
||||
r = run(flow())
|
||||
assert "/status" in r["help"]
|
||||
assert "New videos queued: 1" in r["sync"]
|
||||
assert "Class 01 Aman sir.mp4" in r["check"] and "001, 002" in r["check"]
|
||||
assert r["run"].startswith("Processing started")
|
||||
assert "Processed (uploaded & verified): 1" in r["progress"]
|
||||
assert "Misc• Grammar Practice: 1" in r["folders"]
|
||||
assert "Queue: idle" in r["status"]
|
||||
assert r["failed"] == "No failed items."
|
||||
assert "Unknown command" in r["unknown"]
|
||||
assert r["badcheck"].startswith("Usage")
|
||||
|
||||
|
||||
def test_failed_listing_hides_paths(router, db):
|
||||
run(router.handle(ADMIN, "/sync"))
|
||||
db.transition((CHANNEL_ID, 1), State.FAILED, last_error="cannot open /Users/someone/secret/dir/file.part")
|
||||
reply = run(router.handle(ADMIN, "/failed"))
|
||||
assert "/Users/someone" not in reply and "file.part" in reply
|
||||
@@ -0,0 +1,102 @@
|
||||
from caption_parser import parse_caption
|
||||
from conftest import EXAMPLE_CAPTION
|
||||
|
||||
|
||||
def test_supplied_example():
|
||||
p = parse_caption(EXAMPLE_CAPTION)
|
||||
assert p.indexes == ["001", "002"]
|
||||
assert p.title == "Class 01 Aman sir.mp4"
|
||||
assert p.topic == "Misc• Grammar Practice"
|
||||
assert p.batch == "English ( SSC MAINS )"
|
||||
assert p.extracted_by == "@PUTINxINDIA"
|
||||
assert p.missing == []
|
||||
assert p.warnings == []
|
||||
assert p.raw == EXAMPLE_CAPTION
|
||||
|
||||
|
||||
def test_index_is_never_used_as_title():
|
||||
p = parse_caption("Index: 007\nIndex: 008\nTopic: Nouns")
|
||||
assert p.indexes == ["007", "008"]
|
||||
assert p.title is None
|
||||
assert "title" in p.missing
|
||||
|
||||
|
||||
def test_multiple_index_order_preserved():
|
||||
p = parse_caption("Index: 3\nTitle: A\nIndex: 1\nTopic: T\nIndex: 2")
|
||||
assert p.indexes == ["3", "1", "2"]
|
||||
|
||||
|
||||
def test_fields_in_any_order_and_whitespace_variants():
|
||||
caption = " topic : Error Detection \n\n\nTITLE:Class 02.mp4\r\n batch: X \nExtracted by :@someone\n"
|
||||
p = parse_caption(caption)
|
||||
assert p.topic == "Error Detection"
|
||||
assert p.title == "Class 02.mp4"
|
||||
assert p.batch == "X"
|
||||
assert p.extracted_by == "@someone"
|
||||
|
||||
|
||||
def test_decorations_and_fullwidth_colon():
|
||||
p = parse_caption("🎬 Title： Class 05.mp4\n📚 Topic: Reading Comprehension\n• Batch - ignored\n**Batch:** B1")
|
||||
assert p.title == "Class 05.mp4"
|
||||
assert p.topic == "Reading Comprehension"
|
||||
assert p.batch == "** B1" # value text is preserved verbatim apart from trimming
|
||||
|
||||
|
||||
def test_missing_title_and_topic():
|
||||
p = parse_caption("Index: 001\nBatch: English")
|
||||
assert p.missing == ["title", "topic"]
|
||||
assert "Title field is missing" in p.warnings
|
||||
assert "Topic field is missing" in p.warnings
|
||||
|
||||
|
||||
def test_empty_and_none_caption():
|
||||
for value in (None, "", " \n "):
|
||||
p = parse_caption(value)
|
||||
assert p.missing == ["title", "topic"]
|
||||
assert "caption is empty" in p.warnings
|
||||
|
||||
|
||||
def test_empty_field_values_are_reported():
|
||||
p = parse_caption("Title:\nTopic: \nIndex:")
|
||||
assert p.title is None and p.topic is None
|
||||
assert "Title field is empty" in p.warnings
|
||||
assert "empty Index field ignored" in p.warnings
|
||||
|
||||
|
||||
def test_unicode_preserved():
|
||||
p = parse_caption("Title: कक्षा 01 — व्याकरण (भाग १).mp4\nTopic: हिंदी • व्याकरण\nBatch: Ünïcödé ( 2024 )")
|
||||
assert p.title == "कक्षा 01 — व्याकरण (भाग १).mp4"
|
||||
assert p.topic == "हिंदी • व्याकरण"
|
||||
assert p.batch == "Ünïcödé ( 2024 )"
|
||||
|
||||
|
||||
def test_next_field_is_not_swallowed():
|
||||
p = parse_caption("Title: Class 01\nTopic: Grammar\nBatch: SSC")
|
||||
assert p.title == "Class 01"
|
||||
assert p.topic == "Grammar"
|
||||
|
||||
|
||||
def test_flattened_single_line_caption():
|
||||
p = parse_caption("Index: 001 Title: Class 01 Aman sir.mp4 Topic: Misc• Grammar Practice | Batch: English ( SSC MAINS )")
|
||||
assert p.indexes == ["001"]
|
||||
assert p.title == "Class 01 Aman sir.mp4"
|
||||
assert p.topic == "Misc• Grammar Practice"
|
||||
assert p.batch == "English ( SSC MAINS )"
|
||||
|
||||
|
||||
def test_duplicate_field_keeps_first_and_warns():
|
||||
p = parse_caption("Title: First.mp4\nTitle: Second.mp4\nTopic: T")
|
||||
assert p.title == "First.mp4"
|
||||
assert any("duplicate Title" in w for w in p.warnings)
|
||||
|
||||
|
||||
def test_unknown_fields_and_free_text_kept_in_extra_or_ignored():
|
||||
p = parse_caption("Join: @channel\nSome free text line\nTitle: A.mp4\nTopic: T\nSubtitle: not a title")
|
||||
assert p.title == "A.mp4"
|
||||
assert p.extra["Join"] == "@channel"
|
||||
assert p.extra["Subtitle"] == "not a title"
|
||||
|
||||
|
||||
def test_label_must_be_whole_word():
|
||||
p = parse_caption("Topics: wrong\nTopic: right")
|
||||
assert p.topic == "right"
|
||||
@@ -0,0 +1,140 @@
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
from config import Config, ConfigError, load_json_secret
|
||||
from conftest import make_config
|
||||
from logging_config import RedactSecrets
|
||||
|
||||
|
||||
def _rsa_pem() -> str:
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
|
||||
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
return key.private_bytes(
|
||||
serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()
|
||||
).decode()
|
||||
|
||||
|
||||
SA_INFO = {
|
||||
"type": "service_account",
|
||||
"project_id": "example-project",
|
||||
"private_key_id": "abc123",
|
||||
"client_email": "uploader@example-project.iam.gserviceaccount.com",
|
||||
"client_id": "1",
|
||||
"token_uri": "https://oauth2.googleapis.com/token",
|
||||
}
|
||||
TOKEN_INFO = {
|
||||
"token": "ya29.placeholder-access-token",
|
||||
"refresh_token": "1//placeholder-refresh-token",
|
||||
"client_id": "placeholder.apps.googleusercontent.com",
|
||||
"client_secret": "placeholder-client-secret",
|
||||
"token_uri": "https://oauth2.googleapis.com/token",
|
||||
"expiry": "2999-01-01T00:00:00Z",
|
||||
}
|
||||
|
||||
|
||||
def test_defaults_and_mapped_names(tmp_path):
|
||||
c = make_config(tmp_path, MAX_RETRIES="4", RETRY_DELAY="15", DUPLICATE_POLICY="suffix_message_id",
|
||||
BATCH_SUBFOLDERS_ENABLED="yes", GOOGLE_DRIVE_ROOT_FOLDER_ID="FOLDER123")
|
||||
assert c.max_attempts == 5 and c.retry_base_delay == 15
|
||||
assert c.duplicate_name_policy == "suffix_message_id"
|
||||
assert c.folder_structure == "batch_topic"
|
||||
assert c.drive_root_folder_id == "FOLDER123"
|
||||
assert c.source_channel == "-1001234567890"
|
||||
assert c.fallback_topic == "Uncategorized"
|
||||
assert c.share_mode == "none" # never public by default
|
||||
assert c.drive_scope == "drive.file"
|
||||
|
||||
|
||||
def test_invalid_values_are_all_reported(tmp_path):
|
||||
with pytest.raises(ConfigError) as err:
|
||||
make_config(tmp_path, MAX_RETRIES="many", PROCESS_MODE="sometimes", UPLOAD_CHUNK_SIZE="1000", KEEP_LOCAL_FILES="maybe")
|
||||
text = str(err.value)
|
||||
for name in ("MAX_RETRIES", "PROCESS_MODE", "UPLOAD_CHUNK_SIZE", "KEEP_LOCAL_FILES"):
|
||||
assert name in text
|
||||
|
||||
|
||||
def test_missing_required_telegram_settings():
|
||||
config = Config.from_env({})
|
||||
with pytest.raises(ConfigError) as err:
|
||||
config.validate("telegram")
|
||||
text = str(err.value)
|
||||
assert "TELEGRAM_API_ID" in text and "SOURCE_CHANNEL_ID" in text
|
||||
|
||||
|
||||
def test_heroku_requires_durable_database_and_session_string(tmp_path):
|
||||
c = make_config(tmp_path, DYNO="worker.1", TELEGRAM_MODE="user")
|
||||
problems = c.database_problems() + c.telegram_problems() + c.drive_problems()
|
||||
joined = " ".join(problems)
|
||||
assert "DATABASE_URL" in joined
|
||||
assert "TELEGRAM_SESSION_STRING" in joined
|
||||
assert "GOOGLE_OAUTH_TOKEN_JSON" in joined
|
||||
ok = make_config(
|
||||
tmp_path, DYNO="worker.1", TELEGRAM_MODE="user", TELEGRAM_SESSION_STRING="1AaBbCc-placeholder",
|
||||
DATABASE_URL="postgres://u:p@db.example.com:5432/d", GOOGLE_OAUTH_TOKEN_JSON=json.dumps(TOKEN_INFO),
|
||||
)
|
||||
ok.validate("telegram", "drive", "database")
|
||||
assert ok.database_url == "postgresql://u:p@db.example.com:5432/d?sslmode=require"
|
||||
assert ok.instance_lock_wait > 0
|
||||
|
||||
|
||||
def test_ephemeral_database_must_be_explicit(tmp_path):
|
||||
c = make_config(tmp_path, DYNO="worker.1", ALLOW_EPHEMERAL_DATABASE="true")
|
||||
assert c.database_problems() == []
|
||||
|
||||
|
||||
def test_json_secrets_raw_multiline_and_base64():
|
||||
raw = json.dumps(TOKEN_INFO, indent=2)
|
||||
assert load_json_secret(raw)["refresh_token"] == TOKEN_INFO["refresh_token"]
|
||||
encoded = base64.b64encode(raw.encode()).decode()
|
||||
assert load_json_secret(encoded)["client_id"] == TOKEN_INFO["client_id"]
|
||||
with pytest.raises(ValueError):
|
||||
load_json_secret("{not json")
|
||||
|
||||
|
||||
def test_bad_credential_json_reported(tmp_path):
|
||||
c = make_config(tmp_path, GOOGLE_AUTH_MODE="service_account", GOOGLE_SERVICE_ACCOUNT_JSON='{"type": "service_account"}')
|
||||
assert any("missing keys" in p for p in c.drive_problems())
|
||||
|
||||
|
||||
def test_credentials_load_from_env_without_files(tmp_path):
|
||||
from drive_client import load_credentials
|
||||
|
||||
sa = dict(SA_INFO, private_key=_rsa_pem())
|
||||
c = make_config(tmp_path, GOOGLE_AUTH_MODE="service_account", GOOGLE_SERVICE_ACCOUNT_JSON=json.dumps(sa))
|
||||
assert not c.google_token_file.exists() and not c.google_client_secrets_file.exists()
|
||||
c.validate("drive")
|
||||
creds = load_credentials(c)
|
||||
assert creds.service_account_email == SA_INFO["client_email"]
|
||||
assert creds.scopes == ["https://www.googleapis.com/auth/drive.file"]
|
||||
|
||||
c = make_config(tmp_path, GOOGLE_OAUTH_TOKEN_JSON=base64.b64encode(json.dumps(TOKEN_INFO).encode()).decode())
|
||||
c.validate("drive")
|
||||
creds = load_credentials(c)
|
||||
assert creds.refresh_token == TOKEN_INFO["refresh_token"] and creds.valid
|
||||
assert not c.google_token_file.exists() # nothing written to the ephemeral disk
|
||||
|
||||
|
||||
def test_secrets_are_redacted_from_logs(tmp_path, caplog):
|
||||
sa = dict(SA_INFO, private_key="-----BEGIN PRIVATE KEY-----placeholder")
|
||||
c = make_config(tmp_path, GOOGLE_SERVICE_ACCOUNT_JSON=json.dumps(sa), GOOGLE_OAUTH_TOKEN_JSON=json.dumps(TOKEN_INFO))
|
||||
secrets = c.secrets()
|
||||
assert c.bot_token in secrets and c.api_hash in secrets
|
||||
assert TOKEN_INFO["refresh_token"] in secrets and TOKEN_INFO["client_secret"] in secrets
|
||||
record = logging.LogRecord("x", logging.INFO, __file__, 1, "token=%s refresh=%s", (c.bot_token, TOKEN_INFO["refresh_token"]), None)
|
||||
RedactSecrets(secrets).filter(record)
|
||||
assert c.bot_token not in record.getMessage()
|
||||
assert TOKEN_INFO["refresh_token"] not in record.getMessage()
|
||||
|
||||
|
||||
def test_main_exits_with_clear_error_when_unconfigured(monkeypatch, tmp_path):
|
||||
import main
|
||||
|
||||
monkeypatch.setenv("LOG_DIR", str(tmp_path / "logs"))
|
||||
for name in ("TELEGRAM_API_ID", "TELEGRAM_API_HASH", "TELEGRAM_BOT_TOKEN", "SOURCE_CHANNEL_ID", "DYNO", "DATABASE_URL"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
assert main.main(["check-config"]) == 2
|
||||
@@ -0,0 +1,155 @@
|
||||
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()
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Static deployment-readiness checks (these do NOT deploy anything to Heroku)."""
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
# Top-level keys documented for app.json (https://devcenter.heroku.com/articles/app-json-schema)
|
||||
APP_JSON_KEYS = {
|
||||
"name", "description", "keywords", "website", "repository", "logo", "success_url", "scripts", "env",
|
||||
"formation", "image", "addons", "buildpacks", "environments", "stack",
|
||||
}
|
||||
ENV_KEYS = {"description", "value", "required", "generator"}
|
||||
|
||||
|
||||
def _config_vars() -> set[str]:
|
||||
source = (ROOT / "config.py").read_text()
|
||||
return set(re.findall(r'r\.\w+\(\s*"([A-Z0-9_]+)"', source)) | set(re.findall(r'env\.get\("([A-Z0-9_]+)"\)', source))
|
||||
|
||||
|
||||
def _app_json() -> dict:
|
||||
return json.loads((ROOT / "app.json").read_text())
|
||||
|
||||
|
||||
def test_app_json_structure():
|
||||
manifest = _app_json()
|
||||
assert set(manifest) <= APP_JSON_KEYS
|
||||
assert len(manifest["name"]) <= 30 and manifest["description"]
|
||||
assert manifest["buildpacks"] == [{"url": "heroku/python"}]
|
||||
assert manifest["formation"] == {"worker": {"quantity": 1, "size": "basic"}}
|
||||
assert manifest["addons"][0]["plan"].startswith("heroku-postgresql:")
|
||||
for name, spec in manifest["env"].items():
|
||||
assert set(spec) <= ENV_KEYS, name
|
||||
assert spec["description"], name
|
||||
assert isinstance(spec.get("required", False), bool), name
|
||||
if "value" in spec:
|
||||
assert isinstance(spec["value"], str), name
|
||||
|
||||
|
||||
def test_app_json_only_lists_supported_variables():
|
||||
supported = _config_vars()
|
||||
unsupported = set(_app_json()["env"]) - supported
|
||||
assert not unsupported, f"app.json lists variables the app does not read: {unsupported}"
|
||||
|
||||
|
||||
def test_env_example_documents_every_variable():
|
||||
documented = set(re.findall(r"^([A-Z0-9_]+)=", (ROOT / ".env.example").read_text(), re.MULTILINE))
|
||||
supported = _config_vars() - {"DYNO"} # DYNO is set by Heroku itself
|
||||
assert supported - documented == set()
|
||||
assert documented - supported == set()
|
||||
|
||||
|
||||
def test_procfile_matches_entry_point():
|
||||
procfile = (ROOT / "Procfile").read_text().strip().splitlines()
|
||||
assert procfile == ["worker: python main.py run"]
|
||||
assert set(_app_json()["formation"]) == {line.split(":")[0] for line in procfile}
|
||||
main_source = (ROOT / "main.py").read_text()
|
||||
assert 'sub.add_parser("run"' in main_source and "signal.SIGTERM" in main_source
|
||||
|
||||
|
||||
def test_runtime_and_requirements():
|
||||
assert (ROOT / ".python-version").read_text().strip() == "3.12"
|
||||
assert not (ROOT / "runtime.txt").exists() # deprecated in favour of .python-version
|
||||
reqs = (ROOT / "requirements.txt").read_text()
|
||||
for package in ("Telethon", "google-api-python-client", "google-auth", "google-auth-oauthlib", "psycopg[binary]", "python-dotenv"):
|
||||
assert re.search(rf"^{re.escape(package)}==", reqs, re.MULTILINE), package
|
||||
|
||||
|
||||
def test_gitignore_protects_secrets_and_state():
|
||||
ignored = (ROOT / ".gitignore").read_text().splitlines()
|
||||
for pattern in (".env", "*.session", "credentials/", "token.json", "session_string.txt", "data/", "downloads/", "*.db"):
|
||||
assert pattern in ignored, pattern
|
||||
|
||||
|
||||
def test_readme_deploy_button():
|
||||
readme = (ROOT / "README.md").read_text()
|
||||
assert "https://www.herokucdn.com/deploy/button.svg" in readme
|
||||
assert re.search(r"\]\(https://heroku\.com/deploy\?template=https://github\.com/[^)\s]+\)", readme)
|
||||
@@ -0,0 +1,305 @@
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
|
||||
import httplib2
|
||||
import pytest
|
||||
from googleapiclient.errors import HttpError
|
||||
|
||||
from conftest import FakeDrive
|
||||
from drive_client import (
|
||||
DriveClient,
|
||||
FolderResolver,
|
||||
ResumableUploader,
|
||||
classify_http_error,
|
||||
escape_query,
|
||||
parse_range,
|
||||
)
|
||||
from errors import AuthError, DriveQuotaError, Interrupted, PermanentError, TransientError, UploadSessionExpired
|
||||
|
||||
CHUNK = 256 * 1024
|
||||
|
||||
|
||||
class Resp:
|
||||
def __init__(self, status, headers=None, body=None):
|
||||
self.status_code = status
|
||||
self.headers = headers or {}
|
||||
self._body = body
|
||||
|
||||
def json(self):
|
||||
if self._body is None:
|
||||
raise ValueError("no body")
|
||||
return self._body
|
||||
|
||||
|
||||
class FakeUploadServer:
|
||||
"""Implements Google's resumable upload protocol in memory."""
|
||||
|
||||
def __init__(self):
|
||||
self.sessions = {}
|
||||
self.fail_next_put = [] # exceptions or Resp objects
|
||||
self.initiations = 0
|
||||
self.put_count = 0
|
||||
|
||||
def request(self, method, url, data=None, headers=None, timeout=None):
|
||||
headers = headers or {}
|
||||
if method == "POST":
|
||||
self.initiations += 1
|
||||
uri = f"https://upload/session/{self.initiations}"
|
||||
meta = json.loads(data)
|
||||
self.sessions[uri] = {"data": b"", "size": int(headers["X-Upload-Content-Length"]), "meta": meta, "done": None}
|
||||
return Resp(200, {"Location": uri})
|
||||
session = self.sessions.get(url)
|
||||
if session is None:
|
||||
return Resp(404)
|
||||
content_range = headers["Content-Range"]
|
||||
if content_range.startswith("bytes */"):
|
||||
return self._status(session)
|
||||
self.put_count += 1
|
||||
if self.fail_next_put:
|
||||
failure = self.fail_next_put.pop(0)
|
||||
if isinstance(failure, BaseException):
|
||||
# simulate a connection drop after the server stored half the chunk
|
||||
session["data"] += data[: len(data) // 2]
|
||||
raise failure
|
||||
return failure
|
||||
start = int(re.match(r"bytes (\d+)-", content_range).group(1))
|
||||
assert start == len(session["data"]), "client must continue exactly where the server is"
|
||||
session["data"] += data
|
||||
return self._status(session)
|
||||
|
||||
def _status(self, session):
|
||||
if len(session["data"]) >= session["size"]:
|
||||
if session["done"] is None:
|
||||
data = session["data"]
|
||||
session["done"] = {
|
||||
"id": "file-1",
|
||||
"name": session["meta"]["name"],
|
||||
"size": str(len(data)),
|
||||
"md5Checksum": hashlib.md5(data).hexdigest(),
|
||||
"parents": session["meta"]["parents"],
|
||||
}
|
||||
return Resp(200, body=session["done"])
|
||||
headers = {"Range": f"bytes=0-{len(session['data']) - 1}"} if session["data"] else {}
|
||||
return Resp(308, headers)
|
||||
|
||||
|
||||
def _upload(server, path, *, session_uri=None, stop_after=None, retries=3):
|
||||
sessions, progress = [], []
|
||||
calls = {"n": 0}
|
||||
|
||||
def should_stop():
|
||||
calls["n"] += 1
|
||||
return stop_after is not None and calls["n"] > stop_after
|
||||
|
||||
uploader = ResumableUploader(server, chunk_size=CHUNK, retries=retries, sleep=lambda s: None)
|
||||
result = uploader.upload(
|
||||
path,
|
||||
metadata={"name": "Class 01.mp4", "parents": ["F"], "appProperties": {"tg_source": "1:2"}},
|
||||
mime_type="video/mp4",
|
||||
size=path.stat().st_size,
|
||||
session_uri=session_uri,
|
||||
on_session=sessions.append,
|
||||
on_progress=lambda d, t: progress.append(d),
|
||||
should_stop=should_stop,
|
||||
)
|
||||
return result, sessions, progress
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def big_file(tmp_path):
|
||||
path = tmp_path / "video.bin"
|
||||
path.write_bytes(bytes(range(256)) * (CHUNK * 3 // 256 + 100)) # 3+ chunks
|
||||
return path
|
||||
|
||||
|
||||
def test_full_upload_in_chunks(big_file):
|
||||
server = FakeUploadServer()
|
||||
result, sessions, progress = _upload(server, big_file)
|
||||
assert result["size"] == str(big_file.stat().st_size)
|
||||
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
|
||||
assert len(sessions) == 1 and server.put_count == 4
|
||||
assert progress[-1] == big_file.stat().st_size
|
||||
|
||||
|
||||
def test_pause_then_resume_same_session(big_file):
|
||||
server = FakeUploadServer()
|
||||
with pytest.raises(Interrupted):
|
||||
_upload(server, big_file, stop_after=2)
|
||||
uri = next(iter(server.sessions))
|
||||
stored = len(server.sessions[uri]["data"])
|
||||
assert 0 < stored < big_file.stat().st_size
|
||||
result, sessions, _ = _upload(server, big_file, session_uri=uri)
|
||||
assert sessions == [] # no new session was started
|
||||
assert server.initiations == 1
|
||||
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
|
||||
|
||||
|
||||
def test_network_error_mid_chunk_resumes_from_server_offset(big_file):
|
||||
server = FakeUploadServer()
|
||||
server.fail_next_put = [ConnectionError("reset")]
|
||||
result, _, _ = _upload(server, big_file)
|
||||
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
|
||||
|
||||
|
||||
def test_server_errors_are_retried_then_raised(big_file):
|
||||
server = FakeUploadServer()
|
||||
server.fail_next_put = [Resp(503), Resp(500)]
|
||||
result, _, _ = _upload(server, big_file)
|
||||
assert result["id"] == "file-1"
|
||||
server = FakeUploadServer()
|
||||
server.fail_next_put = [Resp(503)] * 10
|
||||
with pytest.raises(TransientError):
|
||||
_upload(server, big_file, retries=2)
|
||||
|
||||
|
||||
def test_expired_session(big_file):
|
||||
server = FakeUploadServer()
|
||||
# unknown saved session -> a new session is created transparently
|
||||
result, sessions, _ = _upload(server, big_file, session_uri="https://upload/session/expired")
|
||||
assert len(sessions) == 1 and result["id"] == "file-1"
|
||||
server = FakeUploadServer()
|
||||
server.fail_next_put = [Resp(404)]
|
||||
with pytest.raises(UploadSessionExpired):
|
||||
_upload(server, big_file)
|
||||
|
||||
|
||||
def test_quota_and_permission_errors_are_not_retried(big_file):
|
||||
server = FakeUploadServer()
|
||||
server.fail_next_put = [Resp(403, body={"error": {"errors": [{"reason": "storageQuotaExceeded"}], "message": "full"}})]
|
||||
with pytest.raises(DriveQuotaError):
|
||||
_upload(server, big_file)
|
||||
server = FakeUploadServer()
|
||||
server.fail_next_put = [Resp(403, body={"error": {"errors": [{"reason": "insufficientFilePermissions"}]}})]
|
||||
with pytest.raises(PermanentError):
|
||||
_upload(server, big_file)
|
||||
|
||||
|
||||
def test_chunk_size_must_be_aligned():
|
||||
with pytest.raises(ValueError):
|
||||
ResumableUploader(None, chunk_size=1000)
|
||||
|
||||
|
||||
def test_helpers():
|
||||
assert escape_query("It's a \\ test") == "It\\'s a \\\\ test"
|
||||
assert parse_range("bytes=0-1048575") == 1048576
|
||||
assert parse_range(None) == 0
|
||||
assert isinstance(classify_http_error(401), AuthError)
|
||||
assert isinstance(classify_http_error(429), TransientError)
|
||||
rate = classify_http_error(403, "userRateLimitExceeded")
|
||||
assert isinstance(rate, TransientError) and rate.rate_limited
|
||||
assert isinstance(classify_http_error(404), PermanentError)
|
||||
assert isinstance(classify_http_error(502), TransientError)
|
||||
|
||||
|
||||
# -- DriveClient retry wrapper ---------------------------------------------------------
|
||||
|
||||
|
||||
class FakeRequest:
|
||||
def __init__(self, outcome):
|
||||
self.outcome = outcome
|
||||
|
||||
def execute(self):
|
||||
if isinstance(self.outcome, BaseException):
|
||||
raise self.outcome
|
||||
return self.outcome
|
||||
|
||||
|
||||
def _http_error(status, reason=None):
|
||||
body = json.dumps({"error": {"errors": [{"reason": reason}] if reason else [], "message": "x"}}).encode()
|
||||
return HttpError(httplib2.Response({"status": status}), body)
|
||||
|
||||
|
||||
def test_execute_retries_transient_then_succeeds():
|
||||
outcomes = [_http_error(503), OSError("net down"), {"ok": True}]
|
||||
client = DriveClient(None, None, chunk_size=CHUNK, sleep=lambda s: None)
|
||||
assert client._execute(lambda: FakeRequest(outcomes.pop(0))) == {"ok": True}
|
||||
|
||||
|
||||
def test_execute_does_not_retry_permission_errors():
|
||||
calls = []
|
||||
client = DriveClient(None, None, chunk_size=CHUNK, sleep=lambda s: None)
|
||||
|
||||
def build():
|
||||
calls.append(1)
|
||||
return FakeRequest(_http_error(403, "insufficientPermissions"))
|
||||
|
||||
with pytest.raises(PermanentError):
|
||||
client._execute(build)
|
||||
assert len(calls) == 1
|
||||
assert client._execute(lambda: FakeRequest(_http_error(404)), not_found_ok=True) is None
|
||||
|
||||
|
||||
# -- FolderResolver ----------------------------------------------------------------------
|
||||
|
||||
|
||||
class MemoryFolderDB:
|
||||
def __init__(self):
|
||||
self.folders = {}
|
||||
|
||||
def get_folder(self, parent, key):
|
||||
return self.folders.get((parent, key))
|
||||
|
||||
def put_folder(self, parent, key, name, folder_id):
|
||||
self.folders[(parent, key)] = folder_id
|
||||
|
||||
def delete_folder(self, parent, key):
|
||||
self.folders.pop((parent, key), None)
|
||||
|
||||
|
||||
def _resolver(drive, db=None):
|
||||
return FolderResolver(drive, db or MemoryFolderDB(), "root", sleep=lambda s: None)
|
||||
|
||||
|
||||
def test_folder_created_once_and_reused():
|
||||
drive = FakeDrive()
|
||||
resolver = _resolver(drive)
|
||||
first = resolver.resolve(["Misc• Grammar Practice"])
|
||||
second = resolver.resolve(["Misc• Grammar Practice"])
|
||||
assert first == second and drive.created_folders == 1
|
||||
# a fresh process (empty in-memory verification) still reuses the DB mapping
|
||||
db = resolver._db
|
||||
again = FolderResolver(drive, db, "root").resolve(["Misc• Grammar Practice"])
|
||||
assert again == first and drive.created_folders == 1
|
||||
|
||||
|
||||
def test_existing_drive_folder_is_reused_and_duplicates_pick_oldest():
|
||||
drive = FakeDrive()
|
||||
oldest = drive.add_folder("root", "Error Detection")
|
||||
drive.add_folder("root", "Error Detection")
|
||||
resolver = _resolver(drive)
|
||||
assert resolver.resolve(["Error Detection"]) == oldest.id
|
||||
assert drive.created_folders == 0
|
||||
|
||||
|
||||
def test_deleted_cached_folder_is_recreated():
|
||||
drive = FakeDrive()
|
||||
resolver = _resolver(drive)
|
||||
folder_id = resolver.resolve(["Topic"])
|
||||
drive.files[folder_id].trashed = True
|
||||
fresh = FolderResolver(drive, resolver._db, "root")
|
||||
new_id = fresh.resolve(["Topic"])
|
||||
assert new_id != folder_id and drive.created_folders == 2
|
||||
|
||||
|
||||
def test_lost_create_response_does_not_duplicate_folder():
|
||||
drive = FakeDrive()
|
||||
real_create = drive.create_folder
|
||||
|
||||
def flaky_create(parent, name):
|
||||
real_create(parent, name) # Drive created it...
|
||||
raise TransientError("...but the response was lost")
|
||||
|
||||
drive.create_folder = flaky_create
|
||||
resolver = _resolver(drive)
|
||||
folder_id = resolver.resolve(["Reading Comprehension"])
|
||||
assert len(drive.find_folders("root", "Reading Comprehension")) == 1
|
||||
assert drive.files[folder_id].name == "Reading Comprehension"
|
||||
|
||||
|
||||
def test_nested_batch_topic_path():
|
||||
drive = FakeDrive()
|
||||
resolver = _resolver(drive)
|
||||
leaf = resolver.resolve(["English ( SSC MAINS )", "Misc• Grammar Practice"])
|
||||
batch = drive.find_folders("root", "English ( SSC MAINS )")[0]
|
||||
assert drive.files[leaf].parents == (batch.id,)
|
||||
@@ -0,0 +1,66 @@
|
||||
import pytest
|
||||
|
||||
from link_parser import LinkError, internal_channel_id, marked_channel_id, parse_message_link
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url",
|
||||
[
|
||||
"https://t.me/c/3909337424/14",
|
||||
"https://t.me/c/3909337424/14/",
|
||||
"https://t.me/c/3909337424/14?single",
|
||||
"https://www.t.me/c/3909337424/14",
|
||||
"http://t.me/c/3909337424/14#frag",
|
||||
"t.me/c/3909337424/14",
|
||||
"https://telegram.me/c/3909337424/14",
|
||||
"https://t.me/c/3909337424/7/14", # thread segment
|
||||
" https://t.me/c/3909337424/14 ",
|
||||
],
|
||||
)
|
||||
def test_private_links(url):
|
||||
link = parse_message_link(url)
|
||||
assert link.kind == "private"
|
||||
assert link.channel == "3909337424"
|
||||
assert link.message_id == 14
|
||||
assert link.channel_ref == "-1003909337424"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("url", ["https://t.me/my_channel/14", "https://t.me/s/my_channel/14", "t.me/My_Channel/14?single"])
|
||||
def test_public_links(url):
|
||||
link = parse_message_link(url)
|
||||
assert link.kind == "public" and link.message_id == 14
|
||||
assert link.channel.lower() == "my_channel"
|
||||
assert link.channel_ref.lower() == "my_channel"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text,fragment",
|
||||
[
|
||||
("", "send a Telegram message link"),
|
||||
("hello", "not a Telegram"),
|
||||
("https://example.com/c/1/2", "not a Telegram"),
|
||||
("https://t.me/c/3909337424", "must look like"),
|
||||
("https://t.me/c/abc/14", "must look like"),
|
||||
("https://t.me/c/3909337424/0", "positive"),
|
||||
("https://t.me/c/3909337424/x14", "must look like"),
|
||||
("https://t.me/+AbCdEf123", "Invite links"),
|
||||
("https://t.me/joinchat/AbCdEf", "Invite links"),
|
||||
("https://t.me/my_channel", "not to a message"),
|
||||
("https://t.me/addstickers/pack/1", "Unsupported"),
|
||||
("https://t.me/ab/14", "Unsupported"),
|
||||
("ftp://t.me/c/1/2", "not a Telegram"),
|
||||
("https://t.me/c/3909337424/14 https://t.me/c/3909337424/15", "only the link"),
|
||||
],
|
||||
)
|
||||
def test_rejected_links(text, fragment):
|
||||
with pytest.raises(LinkError) as err:
|
||||
parse_message_link(text)
|
||||
assert fragment.lower() in str(err.value).lower()
|
||||
|
||||
|
||||
def test_marked_id_follows_telegram_convention():
|
||||
# -100 prefix is only a coincidence of 10-digit IDs; the real rule is -(10**12 + id)
|
||||
assert marked_channel_id(3909337424) == -1003909337424
|
||||
assert marked_channel_id(12345) == -1000000012345
|
||||
assert internal_channel_id(-1000000012345) == 12345
|
||||
assert parse_message_link("https://t.me/c/12345/9").channel_ref == "-1000000012345"
|
||||
@@ -0,0 +1,84 @@
|
||||
import pytest
|
||||
|
||||
from conftest import video_bytes
|
||||
from errors import PermanentError
|
||||
from naming import add_suffix, build_filename, clean_component, detect_video_extension, normalize_key, same_family
|
||||
|
||||
|
||||
def _name(title, **kw):
|
||||
defaults = dict(message_id=42, original_filename=None, mime_type="video/mp4", detected_ext=".mp4")
|
||||
defaults.update(kw)
|
||||
return build_filename(title, **defaults)
|
||||
|
||||
|
||||
def test_title_used_verbatim():
|
||||
assert _name("Class 01 Aman sir.mp4") == ("Class 01 Aman sir.mp4", [])
|
||||
|
||||
|
||||
def test_extension_added_when_missing():
|
||||
assert _name("Class 01 Aman sir")[0] == "Class 01 Aman sir.mp4"
|
||||
# '.5' is not a video extension, so it stays part of the name
|
||||
assert _name("Lecture 1.5", detected_ext=".mkv")[0] == "Lecture 1.5.mkv"
|
||||
|
||||
|
||||
def test_extension_from_mime_when_undetected():
|
||||
assert _name("Clip", detected_ext=None, mime_type="video/x-matroska")[0] == "Clip.mkv"
|
||||
|
||||
|
||||
def test_compatible_extension_kept():
|
||||
assert _name("Talk.mov", detected_ext=".mp4") == ("Talk.mov", [])
|
||||
|
||||
|
||||
def test_mismatch_kept_by_default_but_reported():
|
||||
name, warnings = _name("Talk.mp4", detected_ext=".mkv")
|
||||
assert name == "Talk.mp4"
|
||||
assert any("does not match" in w for w in warnings)
|
||||
|
||||
|
||||
def test_mismatch_append_policy():
|
||||
assert _name("Talk.mp4", detected_ext=".mkv", mismatch_policy="append")[0] == "Talk.mp4.mkv"
|
||||
|
||||
|
||||
def test_missing_title_fallbacks():
|
||||
name, warnings = _name(None, original_filename="orig_video.mp4")
|
||||
assert name == "orig_video.mp4" and warnings
|
||||
name, warnings = _name(" ", original_filename=None)
|
||||
assert name == "telegram_42.mp4" and warnings
|
||||
|
||||
|
||||
def test_missing_title_fail_policy():
|
||||
with pytest.raises(PermanentError):
|
||||
_name(None, missing_title_policy="fail")
|
||||
|
||||
|
||||
def test_sanitizing():
|
||||
assert clean_component("a/b\\c\x00d\n e") == "a_b_c d e"
|
||||
assert clean_component("Misc• Grammar Practice") == "Misc• Grammar Practice"
|
||||
assert clean_component('Q: what? "x" <y>|z*', strict=True) == "Q_ what_ _x_ _y__z_"
|
||||
assert len(clean_component("x" * 500 + ".mp4")) == 200
|
||||
assert clean_component("x" * 500 + ".mp4").endswith(".mp4")
|
||||
|
||||
|
||||
def test_add_suffix():
|
||||
assert add_suffix("Class 01 Aman sir.mp4", "(2)") == "Class 01 Aman sir (2).mp4"
|
||||
assert add_suffix("Class 01", "[77]") == "Class 01 [77]"
|
||||
|
||||
|
||||
def test_detect_container(tmp_path):
|
||||
cases = {
|
||||
".mp4": video_bytes(1000),
|
||||
".mkv": b"\x1a\x45\xdf\xa3" + b"\x00" * 30 + b"matroska" + b"\x00" * 100,
|
||||
".webm": b"\x1a\x45\xdf\xa3" + b"\x00" * 30 + b"webm" + b"\x00" * 100,
|
||||
".avi": b"RIFF\x00\x00\x00\x00AVI LIST",
|
||||
None: b"hello world, not a video",
|
||||
}
|
||||
for expected, data in cases.items():
|
||||
path = tmp_path / "f"
|
||||
path.write_bytes(data)
|
||||
assert detect_video_extension(path) == expected
|
||||
|
||||
|
||||
def test_families_and_keys():
|
||||
assert same_family(".MP4", ".mov")
|
||||
assert not same_family(".mp4", ".mkv")
|
||||
assert normalize_key(" Misc• Grammar Practice ") == normalize_key("misc• grammar practice")
|
||||
@@ -0,0 +1,107 @@
|
||||
import asyncio
|
||||
|
||||
from telethon import errors
|
||||
from telethon.tl import types
|
||||
|
||||
from conftest import RecordingNotifier
|
||||
from errors import PermanentError, TransientError
|
||||
from progress import ProgressReporter, TransferProgress, format_bytes, format_duration
|
||||
from telegram_client import classify_media, map_telegram_error, parse_channel_reference, to_source_message
|
||||
|
||||
|
||||
# -- progress ------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_formatting():
|
||||
assert format_bytes(512) == "512 B"
|
||||
assert format_bytes(1536) == "1.5 KB"
|
||||
assert format_bytes(5 * 1024**3) == "5.0 GB"
|
||||
assert format_duration(42) == "42s"
|
||||
assert format_duration(125) == "2m 5s"
|
||||
assert format_duration(3 * 3600 + 120) == "3h 2m"
|
||||
|
||||
|
||||
def test_speed_and_eta():
|
||||
now = [0.0]
|
||||
p = TransferProgress(total=100 * 1024 * 1024, clock=lambda: now[0])
|
||||
p.update(0)
|
||||
now[0] = 10.0
|
||||
p.update(50 * 1024 * 1024)
|
||||
assert round(p.speed) == 5 * 1024 * 1024
|
||||
assert round(p.eta) == 10
|
||||
assert p.percent == 50
|
||||
text = p.describe()
|
||||
assert "50%" in text and "MB/s" in text and "ETA 10s" in text
|
||||
|
||||
|
||||
def test_reporter_edits_single_message():
|
||||
notifier = RecordingNotifier()
|
||||
reporter = ProgressReporter(notifier, [1, 2], interval=60)
|
||||
|
||||
async def go():
|
||||
reporter.board.begin_item(12, 150, "Misc• Grammar Practice", "Class 01 Aman sir.mp4")
|
||||
reporter.board.start_download(1000)
|
||||
reporter.board.download.update(640)
|
||||
await reporter.refresh()
|
||||
await reporter.refresh() # unchanged -> nothing sent
|
||||
reporter.board.download.update(900)
|
||||
await reporter.refresh()
|
||||
|
||||
asyncio.run(go())
|
||||
assert len(notifier.sent) == 2 # one per chat
|
||||
assert len(notifier.edits) == 2 # then edited
|
||||
first = notifier.sent[0][1]
|
||||
assert first.splitlines()[0] == "Processing video 12 of 150"
|
||||
assert "Topic: Misc• Grammar Practice" in first and "Download: 64%" in first and "Upload: Waiting" in first
|
||||
|
||||
|
||||
# -- telegram media classification ----------------------------------------------------------
|
||||
|
||||
|
||||
def _doc(attributes, mime="video/mp4"):
|
||||
return types.Document(id=77, access_hash=1, file_reference=b"", date=None, mime_type=mime, size=1234, dc_id=1,
|
||||
attributes=attributes)
|
||||
|
||||
|
||||
def _msg(media, text="Title: A.mp4"):
|
||||
return types.Message(id=5, peer_id=types.PeerChannel(1), date=None, message=text, media=media)
|
||||
|
||||
|
||||
def test_classify_video_message():
|
||||
video = types.DocumentAttributeVideo(duration=60, w=1280, h=720)
|
||||
info = classify_media(_msg(types.MessageMediaDocument(document=_doc([video]))), include_documents=False)
|
||||
assert info.kind == "video" and info.size == 1234 and info.media_id == 77
|
||||
sm = to_source_message(_msg(types.MessageMediaDocument(document=_doc([video]))), -100, False)
|
||||
assert sm.caption == "Title: A.mp4" and sm.message_id == 5
|
||||
|
||||
|
||||
def test_classify_document_video_only_when_enabled():
|
||||
doc = _doc([types.DocumentAttributeFilename("lecture.mkv")], mime="application/octet-stream")
|
||||
message = _msg(types.MessageMediaDocument(document=doc))
|
||||
assert classify_media(message, include_documents=False) is None
|
||||
assert classify_media(message, include_documents=True).file_name == "lecture.mkv"
|
||||
|
||||
|
||||
def test_unsupported_media_is_ignored():
|
||||
gif = _doc([types.DocumentAttributeVideo(duration=1, w=1, h=1), types.DocumentAttributeAnimated()])
|
||||
note = _doc([types.DocumentAttributeVideo(duration=1, w=1, h=1, round_message=True)])
|
||||
pdf = _doc([types.DocumentAttributeFilename("notes.pdf")], mime="application/pdf")
|
||||
for doc in (gif, note, pdf):
|
||||
assert classify_media(_msg(types.MessageMediaDocument(document=doc)), include_documents=True) is None
|
||||
photo = types.MessageMediaPhoto(photo=types.PhotoEmpty(id=1))
|
||||
assert classify_media(_msg(photo), include_documents=True) is None
|
||||
assert classify_media(_msg(None, text="text only"), include_documents=True) is None
|
||||
|
||||
|
||||
def test_channel_reference_parsing():
|
||||
assert parse_channel_reference("-1001234567890") == -1001234567890
|
||||
assert parse_channel_reference("@my_channel") == "my_channel"
|
||||
assert parse_channel_reference("https://t.me/my_channel") == "my_channel"
|
||||
assert parse_channel_reference("https://t.me/c/1234567890/55") == -1001234567890
|
||||
|
||||
|
||||
def test_telegram_error_mapping():
|
||||
flood = map_telegram_error(errors.FloodWaitError(request=None, capture=42))
|
||||
assert isinstance(flood, TransientError) and flood.retry_after == 42 and flood.rate_limited
|
||||
assert isinstance(map_telegram_error(errors.ChannelPrivateError(request=None)), PermanentError)
|
||||
assert isinstance(map_telegram_error(ConnectionError("x")), TransientError)
|
||||
@@ -0,0 +1,438 @@
|
||||
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]
|
||||
@@ -0,0 +1,447 @@
|
||||
"""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())
|
||||
Reference in New Issue
Block a user