"""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)