374 lines
13 KiB
Python
374 lines
13 KiB
Python
"""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)
|