fix(stability): стадия 3 — раздача через FileResponse, чистый ASGI-логгер, больший executor #11
+210
-137
@@ -15,7 +15,6 @@ import re
|
||||
import uuid
|
||||
import mimetypes
|
||||
from typing import List, Union, Any
|
||||
from urllib.parse import quote
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
@@ -23,9 +22,9 @@ import time
|
||||
from contextlib import asynccontextmanager
|
||||
import random
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import uvloop
|
||||
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.background import BackgroundTask
|
||||
import sys
|
||||
|
||||
@@ -34,14 +33,15 @@ from pyrogram import errors
|
||||
from pyrogram.types import Message
|
||||
from pyrogram.enums import MessageMediaType
|
||||
from fastapi import FastAPI, HTTPException, Response, Request
|
||||
from fastapi.responses import HTMLResponse, FileResponse, StreamingResponse
|
||||
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
|
||||
from telegram_client import TelegramClient
|
||||
from config import get_settings, setup_logging
|
||||
from rss_generator import generate_channel_rss, generate_channel_html
|
||||
from post_parser import PostParser
|
||||
from url_signer import verify_media_digest, generate_media_digest
|
||||
from file_io import (DB_PATH, init_db_sync, get_all_media_file_ids_sync,
|
||||
update_media_file_access_sync, remove_media_file_ids_sync,
|
||||
update_media_file_access_sync, update_media_file_access_bulk_sync,
|
||||
remove_media_file_ids_sync,
|
||||
get_mime_type_sync, set_mime_type_sync)
|
||||
|
||||
# Global python-magic instance for MIME type detection
|
||||
@@ -51,14 +51,35 @@ magic_mime = magic.Magic(mime=True)
|
||||
class ZeroSizeFileError(Exception):
|
||||
"""Custom exception for zero-size files found or downloaded."""
|
||||
|
||||
class RequestLoggingMiddleware(BaseHTTPMiddleware):
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
# Log only method and URL at debug level to avoid flooding logs on active RSS polling
|
||||
logger.debug(f"Request: {request.method} {request.url}")
|
||||
class RequestLoggingMiddleware:
|
||||
"""Pure-ASGI request logger (no BaseHTTPMiddleware).
|
||||
|
||||
BaseHTTPMiddleware runs the downstream app in a separate anyio task and pumps the
|
||||
response through an in-memory stream, which adds per-request overhead and interacts
|
||||
badly with streaming bodies, background tasks and client cancellation. This plain
|
||||
ASGI middleware only wraps `send` to observe the response status line, so it never
|
||||
buffers the body — the FileResponse stream flows straight through untouched.
|
||||
"""
|
||||
def __init__(self, app):
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
# Log only method and path (with query, matching the old request.url logging) at
|
||||
# debug level to avoid flooding logs on active RSS polling.
|
||||
_qs = scope.get("query_string") or b""
|
||||
_path = scope["path"] + (f"?{_qs.decode('latin-1')}" if _qs else "")
|
||||
logger.debug(f"Request: {scope['method']} {_path}")
|
||||
|
||||
async def send_wrapper(message):
|
||||
if message["type"] == "http.response.start":
|
||||
logger.debug(f"Response status: {message['status']}")
|
||||
await send(message)
|
||||
|
||||
try:
|
||||
response = await call_next(request)
|
||||
logger.debug(f"Response status: {response.status_code}")
|
||||
return response
|
||||
await self.app(scope, receive, send_wrapper)
|
||||
except Exception as e:
|
||||
logger.error(f"Request processing error: {str(e)}")
|
||||
raise
|
||||
@@ -71,6 +92,12 @@ Config = get_settings()
|
||||
HTTP_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(3) # semaphore for live HTTP media requests
|
||||
BACKGROUND_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(2) # semaphore for background cache worker
|
||||
download_queue = asyncio.Queue(maxsize=100)
|
||||
# How stale a temp_* file's mtime must be before a serve refreshes it (keeps the 1h
|
||||
# sweeper from deleting an actively-viewed video). Well below 1h so the file stays alive,
|
||||
# but large enough that the mtime — and thus FileResponse's ETag — is stable within any
|
||||
# such window; a view running longer than one interval costs at most one safe 200
|
||||
# If-Range restart per interval (a full re-fetch, never corruption).
|
||||
TEMP_MTIME_REFRESH_INTERVAL = 300 # seconds
|
||||
|
||||
# In-flight download dedup registry: maps (channel, post_id, file_unique_id) to the
|
||||
# shared Future of an ongoing download. The FIRST request for a key runs the download in
|
||||
@@ -113,11 +140,70 @@ async def _supervised(factory, name: str, min_restart_interval: float = 60.0):
|
||||
await asyncio.sleep(min_restart_interval - elapsed)
|
||||
|
||||
|
||||
# Access-time write accumulator. A /media cache hit used to touch SQLite on the hot path
|
||||
# (a threadpool hop + connect + UPDATE per request), which starves the threadpool under
|
||||
# active RSS polling. Instead a cache hit just records the timestamp here — a dict write
|
||||
# on the single-threaded event loop is cheap and atomic — and a periodic background task
|
||||
# flushes the whole batch to SQLite in one executemany. Keys use str(channel) to stay
|
||||
# consistent with the string form written at insert time and to not lean on SQLite's
|
||||
# implicit column-affinity coercion (the channel column is TEXT, so a bound int would be
|
||||
# coerced and still match — but we key the accumulator by the same type we store, rather
|
||||
# than depend on that).
|
||||
ACCESS_FLUSH_INTERVAL = 60 # seconds between access-time flushes
|
||||
_access_updates: dict[tuple[str, int, str], float] = {}
|
||||
|
||||
|
||||
async def _flush_access_updates() -> None:
|
||||
"""Flush the accumulated access timestamps to SQLite in one bulk UPDATE.
|
||||
|
||||
Snapshot-then-clear atomically on the loop: capture the current dict reference and
|
||||
replace the module global with a fresh empty dict in ONE synchronous step (before any
|
||||
await), so cache-hit writes arriving DURING the flush land in the new dict and are not
|
||||
lost. The bulk UPDATE runs off-loop via asyncio.to_thread. An empty batch is a no-op.
|
||||
"""
|
||||
global _access_updates
|
||||
if not _access_updates:
|
||||
return
|
||||
pending = _access_updates
|
||||
_access_updates = {}
|
||||
entries = [(channel, post_id, file_unique_id, added)
|
||||
for (channel, post_id, file_unique_id), added in pending.items()]
|
||||
try:
|
||||
await asyncio.to_thread(update_media_file_access_bulk_sync, DB_PATH, entries)
|
||||
except Exception:
|
||||
# Bulk write failed: re-queue this batch so the access-times are not lost (a lost
|
||||
# timestamp would eventually evict a still-used file from the 20-day cache). Use
|
||||
# setdefault so any FRESHER write accumulated during the flush is never overwritten
|
||||
# by our stale snapshot. Runs on the loop with no await before the mutation, so this
|
||||
# is race-free. Re-raise so the flush loop logs it.
|
||||
for key, added in pending.items():
|
||||
_access_updates.setdefault(key, added)
|
||||
raise
|
||||
|
||||
|
||||
async def _access_flush_loop() -> None:
|
||||
"""Periodically flush the access-time accumulator (runs under _supervised)."""
|
||||
while True:
|
||||
await asyncio.sleep(ACCESS_FLUSH_INTERVAL)
|
||||
try:
|
||||
await _flush_access_updates()
|
||||
except Exception as e:
|
||||
# Log and keep looping: a transient SQLite error must not drop the batch's
|
||||
# successors. (_supervised still restarts us if this ever raises out.)
|
||||
logger.error(f"access_flush_error: {e}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI):
|
||||
setup_logging(Config["log_level"])
|
||||
|
||||
|
||||
# Enlarge the default threadpool: SQLite/python-magic/pickle/os.walk all run via
|
||||
# asyncio.to_thread, and the interpreter default (min(32, cpu+4) = 5-6 on a 1-2 CPU
|
||||
# container) is too small under load. Configurable via IO_THREAD_POOL_SIZE.
|
||||
loop = asyncio.get_running_loop()
|
||||
io_executor = ThreadPoolExecutor(max_workers=Config["io_thread_pool_size"], thread_name_prefix="io")
|
||||
loop.set_default_executor(io_executor)
|
||||
|
||||
base_cache_dir = os.path.abspath("./data/cache")
|
||||
os.makedirs(base_cache_dir, exist_ok=True) # Create cache directory
|
||||
|
||||
@@ -129,9 +215,11 @@ async def lifespan(_: FastAPI):
|
||||
# CRITICAL and restarted, so a crash can no longer silently stop cache sweeping or downloads.
|
||||
background_task = asyncio.create_task(_supervised(cache_media_files, "cache_media_files"))
|
||||
worker_task = asyncio.create_task(_supervised(background_download_worker, "background_download_worker"))
|
||||
access_flush_task = asyncio.create_task(_supervised(_access_flush_loop, "access_flush_loop"))
|
||||
yield
|
||||
background_task.cancel() # Cleanup
|
||||
worker_task.cancel()
|
||||
access_flush_task.cancel()
|
||||
try:
|
||||
await background_task
|
||||
except asyncio.CancelledError:
|
||||
@@ -140,7 +228,20 @@ async def lifespan(_: FastAPI):
|
||||
await worker_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
try:
|
||||
await access_flush_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
# Final flush AFTER the loop task is cancelled (no race with a loop-driven flush) and
|
||||
# BEFORE the threadpool is shut down (to_thread still has its executor), so the last
|
||||
# <=ACCESS_FLUSH_INTERVAL seconds of access-times are persisted on shutdown.
|
||||
try:
|
||||
await _flush_access_updates()
|
||||
except Exception as e:
|
||||
logger.error(f"access_flush_shutdown_error: {e}")
|
||||
await client.stop()
|
||||
# Shut the io threadpool down so its threads don't linger past a reload/restart.
|
||||
io_executor.shutdown(wait=False)
|
||||
|
||||
app = FastAPI(title="Pyrogram Bridge", lifespan=lifespan)
|
||||
app.add_middleware(RequestLoggingMiddleware)
|
||||
@@ -244,19 +345,45 @@ async def delayed_delete_file(file_path: str, delay: int = 300) -> None:
|
||||
|
||||
|
||||
async def prepare_file_response(file_path: str, request: Request, delete_after: bool = False,
|
||||
media_key: tuple[str, int, str] | None = None) -> StreamingResponse:
|
||||
"""Prepare a streaming file response with HTTP Range request support."""
|
||||
if not os.path.exists(file_path):
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
media_key: tuple[str, int, str] | None = None) -> Response:
|
||||
"""Serve a cached media file via Starlette's FileResponse.
|
||||
|
||||
FileResponse handles Range/If-Range/206/416/multipart and sets
|
||||
Accept-Ranges/ETag/Last-Modified itself, and reads the file efficiently (no per-64KB
|
||||
to_thread hop that starved the threadpool). We keep: the early 404 pre-check, the MIME
|
||||
logic (python-magic + SQLite type cache), the stage-2 temp_* mtime touch, and the
|
||||
delete_after BackgroundTask.
|
||||
"""
|
||||
# `request` is unused now that FileResponse parses the Range header itself, but the
|
||||
# signature is kept for call-site compatibility (and future needs).
|
||||
|
||||
# Keep an actively-viewed large-video temp file alive: refresh its mtime so the 1h
|
||||
# sweeper (which deletes temp_* by mtime) won't remove it out from under a viewer.
|
||||
# DEBOUNCED: FileResponse derives ETag/Last-Modified from mtime, so touching on EVERY
|
||||
# serve would change the validators per request. We only refresh when the mtime is
|
||||
# already older than TEMP_MTIME_REFRESH_INTERVAL — far below the 1h sweeper window, so
|
||||
# the file stays alive, and the ETag is stable within any such window (a view running
|
||||
# longer than one interval costs at most one safe 200 If-Range restart per interval).
|
||||
if os.path.basename(file_path).startswith("temp_"):
|
||||
try:
|
||||
await asyncio.to_thread(os.utime, file_path, None)
|
||||
age = time.time() - await asyncio.to_thread(os.path.getmtime, file_path)
|
||||
if age > TEMP_MTIME_REFRESH_INTERVAL:
|
||||
await asyncio.to_thread(os.utime, file_path, None)
|
||||
except OSError as e:
|
||||
logger.debug(f"Failed to refresh mtime for {file_path}: {e}")
|
||||
|
||||
# Take ONE authoritative stat and hand it to FileResponse as stat_result. This both
|
||||
# (a) preserves the 404 semantics: FileResponse with stat_result=None re-stats at
|
||||
# send-time and raises a RuntimeError (-> 500, escaping this handler's try/except) if
|
||||
# the file was swept between here and the send; and (b) makes the ETag/Last-Modified
|
||||
# reflect exactly the mtime observed after the optional touch above. The remaining
|
||||
# narrow window (deleted between this stat and FileResponse's own open) truncates the
|
||||
# body — that pre-existed the FileResponse migration and is not handled here.
|
||||
try:
|
||||
stat_result = await asyncio.to_thread(os.stat, file_path)
|
||||
except FileNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
|
||||
media_type: str | None = None
|
||||
|
||||
if media_key is not None:
|
||||
@@ -281,107 +408,26 @@ async def prepare_file_response(file_path: str, request: Request, delete_after:
|
||||
|
||||
logger.debug(f"Determined media type for {os.path.basename(file_path)}: {media_type}")
|
||||
|
||||
try:
|
||||
total_size = os.path.getsize(file_path)
|
||||
except FileNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="File not found")
|
||||
range_header = request.headers.get("range")
|
||||
|
||||
if range_header:
|
||||
# Parse Range header according to RFC 7233
|
||||
try:
|
||||
range_value = range_header.strip()
|
||||
if not range_value.startswith("bytes="):
|
||||
raise ValueError("Only bytes ranges are supported")
|
||||
range_spec = range_value[len("bytes="):]
|
||||
if range_spec.startswith("-"):
|
||||
# Suffix range: bytes=-N (last N bytes)
|
||||
suffix_length = int(range_spec[1:])
|
||||
start = max(0, total_size - suffix_length)
|
||||
end = total_size - 1
|
||||
elif range_spec.endswith("-"):
|
||||
# Open-ended range: bytes=START-
|
||||
start = int(range_spec[:-1])
|
||||
end = total_size - 1
|
||||
else:
|
||||
# Full range: bytes=START-END
|
||||
start_str, end_str = range_spec.split("-", 1)
|
||||
start = int(start_str)
|
||||
end = int(end_str)
|
||||
except (ValueError, IndexError) as e:
|
||||
logger.debug(f"Invalid Range header '{range_header}': {e}")
|
||||
return Response(
|
||||
status_code=416,
|
||||
headers={"Content-Range": f"bytes */{total_size}"}
|
||||
)
|
||||
|
||||
# Clamp end to file size - 1 (RFC 7233 allows end >= total_size)
|
||||
end = min(end, total_size - 1)
|
||||
|
||||
# If start is beyond file size, return 416
|
||||
if start >= total_size:
|
||||
return Response(
|
||||
status_code=416,
|
||||
headers={"Content-Range": f"bytes */{total_size}"}
|
||||
)
|
||||
|
||||
content_length = end - start + 1
|
||||
status_code = 206
|
||||
headers = {
|
||||
"Content-Disposition": (
|
||||
f"inline; filename=\"{os.path.basename(file_path)}\"; "
|
||||
f"filename*=UTF-8''{quote(os.path.basename(file_path))}"
|
||||
),
|
||||
# Files are addressed by file_unique_id which is immutable in Telegram,
|
||||
# so it is safe to cache them aggressively on the client side.
|
||||
"Cache-Control": "public, max-age=86400, immutable",
|
||||
"Accept-Ranges": "bytes",
|
||||
"Content-Range": f"bytes {start}-{end}/{total_size}",
|
||||
"Content-Length": str(content_length),
|
||||
}
|
||||
else:
|
||||
# No Range header — serve full file with status 200
|
||||
start = 0
|
||||
end = total_size - 1
|
||||
status_code = 200
|
||||
headers = {
|
||||
"Content-Disposition": (
|
||||
f"inline; filename=\"{os.path.basename(file_path)}\"; "
|
||||
f"filename*=UTF-8''{quote(os.path.basename(file_path))}"
|
||||
),
|
||||
# Files are addressed by file_unique_id which is immutable in Telegram,
|
||||
# so it is safe to cache them aggressively on the client side.
|
||||
"Cache-Control": "public, max-age=86400, immutable",
|
||||
"Accept-Ranges": "bytes",
|
||||
"Content-Length": str(total_size),
|
||||
}
|
||||
|
||||
chunk_size = 64 * 1024 # 64 KB chunks
|
||||
|
||||
async def file_chunk_generator():
|
||||
"""Async generator that reads file in chunks; each chunk is a separate open/seek/read/close."""
|
||||
bytes_remaining = end - start + 1
|
||||
offset = start
|
||||
while bytes_remaining > 0:
|
||||
to_read = min(chunk_size, bytes_remaining)
|
||||
def read_at(path, off, size):
|
||||
# Open, seek, read, and close within a single thread call to avoid fd leaks
|
||||
with open(path, "rb") as f:
|
||||
f.seek(off)
|
||||
return f.read(size)
|
||||
data = await asyncio.to_thread(read_at, file_path, offset, to_read)
|
||||
if not data:
|
||||
break
|
||||
bytes_remaining -= len(data)
|
||||
offset += len(data)
|
||||
yield data
|
||||
|
||||
# Delete the temporary file once the response has been fully sent (stage-2 delete_after).
|
||||
# FileResponse runs this BackgroundTask after streaming the body.
|
||||
background = BackgroundTask(delayed_delete_file, file_path) if delete_after else None
|
||||
return StreamingResponse(
|
||||
content=file_chunk_generator(),
|
||||
status_code=status_code,
|
||||
|
||||
# FileResponse handles Range/If-Range/206/416/multipart and sets
|
||||
# Accept-Ranges/ETag/Last-Modified itself (from the stat_result we pass). Do NOT
|
||||
# hand-build Content-Disposition: FileResponse forms it from filename= —
|
||||
# `inline; filename="x"` for an ASCII name, adding `filename*=UTF-8''x` only for a
|
||||
# non-ASCII name. It uses setdefault, so a manual header would OVERRIDE it, not double
|
||||
# it; letting FileResponse own it keeps the RFC 5987 encoding correct.
|
||||
#
|
||||
# Files are addressed by file_unique_id which is immutable in Telegram, so it is safe
|
||||
# to cache them aggressively on the client side.
|
||||
return FileResponse(
|
||||
file_path,
|
||||
media_type=media_type,
|
||||
headers=headers,
|
||||
filename=os.path.basename(file_path),
|
||||
content_disposition_type="inline",
|
||||
stat_result=stat_result,
|
||||
headers={"Cache-Control": "public, max-age=86400, immutable"},
|
||||
background=background,
|
||||
)
|
||||
|
||||
@@ -559,16 +605,11 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu
|
||||
logger.error(f"cleanup_error: Failed to remove zero-size cached file {cache_path}: {e}")
|
||||
# Do not raise error here, proceed to download below
|
||||
else:
|
||||
# File exists and is not zero size, update access timestamp and return
|
||||
# File exists and is not zero size, record access timestamp and return.
|
||||
# Record into the accumulator instead of touching SQLite on the hot path; the
|
||||
# background flush persists it. Key channel as str(channel) — see _access_updates.
|
||||
logger.info(f"Found cached media file: {cache_path}")
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
update_media_file_access_sync,
|
||||
DB_PATH, str(channel), post_id, file_unique_id,
|
||||
datetime.now().timestamp()
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update timestamp for {channel}/{post_id}/{file_unique_id}: {str(e)}")
|
||||
_access_updates[(str(channel), post_id, file_unique_id)] = datetime.now().timestamp()
|
||||
return cache_path, False
|
||||
|
||||
file_id = await find_file_id_in_message(message, file_unique_id)
|
||||
@@ -1024,6 +1065,46 @@ async def get_raw_post_json(channel: str, post_id: int, request: Request, token:
|
||||
raise HTTPException(status_code=500, detail=error_message) from e
|
||||
|
||||
|
||||
@app.get("/ping")
|
||||
async def ping() -> JSONResponse:
|
||||
"""Lightweight liveness probe for the container healthcheck.
|
||||
|
||||
Reflects process/event-loop liveness (always answers in microseconds) and TG liveness
|
||||
from the watchdog's last-probe data. It MUST NOT issue any Telegram RPC (no get_me,
|
||||
no safe_get_messages), touch SQLite, or walk the filesystem — that is the whole point:
|
||||
it stays instant and truthful even while a real TG RPC is hung. It only reads the
|
||||
already-recorded watchdog timestamp and the is_connected bool.
|
||||
"""
|
||||
age = client.watchdog_last_ok_age() # seconds since last OK probe, None if never
|
||||
# is_connected is None before client.start() and a bool afterwards; coerce so the JSON
|
||||
# "connected" field is always a bool (never null) and the pre-start window reports false.
|
||||
connected = bool(client.client.is_connected)
|
||||
threshold = Config["tg_ping_unhealthy_after"]
|
||||
# age is None right after boot: the watchdog hasn't run its first probe yet. Treat that
|
||||
# as healthy (gate on connected only) so a freshly-started container is not killed before
|
||||
# its first probe cycle — otherwise start_period would have to cover a full watchdog interval.
|
||||
#
|
||||
# The staleness branch (age >= threshold => degraded) is only meaningful while the watchdog
|
||||
# is running to refresh age. With the watchdog DISABLED (TG_WATCHDOG_ENABLED=false) nothing
|
||||
# refreshes age — yet a disconnect-flap restart can still stamp it once (see _restart_client,
|
||||
# which runs before the watchdog-enabled gate), after which age only grows. Letting that
|
||||
# stale age drive /ping to 503 would spuriously fail the container healthcheck on a live
|
||||
# connection and trigger an autoheal restart. So gate staleness on the watchdog being on;
|
||||
# with it off, /ping is a pure connectivity check (no zombie-session detection — that
|
||||
# TG-liveness signal only exists while the watchdog runs).
|
||||
healthy = connected and (
|
||||
not Config["tg_watchdog_enabled"] or age is None or age < threshold
|
||||
)
|
||||
return JSONResponse(
|
||||
{
|
||||
"status": "ok" if healthy else "degraded",
|
||||
"connected": connected,
|
||||
"last_probe_age_s": round(age, 1) if age is not None else None,
|
||||
"threshold_s": threshold,
|
||||
},
|
||||
status_code=200 if healthy else 503,
|
||||
)
|
||||
|
||||
@app.get("/health")
|
||||
@app.get("/health/{token}")
|
||||
async def health_check(request: Request, token: str | None = None) -> Response:
|
||||
@@ -1069,7 +1150,7 @@ async def health_check(request: Request, token: str | None = None) -> Response:
|
||||
|
||||
@app.get("/media/{channel}/{post_id}/{file_unique_id}/{digest}", response_model=None)
|
||||
@app.get("/media/{channel}/{post_id}/{file_unique_id}", response_model=None)
|
||||
async def get_media(channel: str, post_id: int, file_unique_id: str, request: Request, digest: str | None = None) -> StreamingResponse|Response:
|
||||
async def get_media(channel: str, post_id: int, file_unique_id: str, request: Request, digest: str | None = None) -> Response:
|
||||
try:
|
||||
url = f"{channel}/{post_id}/{file_unique_id}"
|
||||
if not verify_media_digest(url, digest):
|
||||
@@ -1095,17 +1176,9 @@ async def get_media(channel: str, post_id: int, file_unique_id: str, request: Re
|
||||
if os.path.exists(cache_path) and os.path.getsize(cache_path) > 0:
|
||||
# File is already in cache — skip semaphore and serve directly
|
||||
logger.info(f"pre_semaphore_cache_hit: {channel}/{post_id}/{file_unique_id}")
|
||||
# Fire-and-forget timestamp update with error handling to avoid silent failures
|
||||
async def _update_access(_ch, _pid, _fid):
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
update_media_file_access_sync,
|
||||
DB_PATH, str(_ch), _pid, _fid,
|
||||
datetime.now().timestamp()
|
||||
)
|
||||
except Exception as _e:
|
||||
logger.warning(f"Failed to update access time for {_ch}/{_pid}/{_fid}: {_e}")
|
||||
asyncio.create_task(_update_access(channel, post_id, file_unique_id))
|
||||
# Record the access time into the accumulator instead of firing a per-hit
|
||||
# SQLite write. Key channel as str(channel) — see _access_updates.
|
||||
_access_updates[(str(channel), post_id, file_unique_id)] = datetime.now().timestamp()
|
||||
return await prepare_file_response(cache_path, request=request,
|
||||
media_key=(str(channel), post_id, file_unique_id))
|
||||
|
||||
|
||||
@@ -123,10 +123,25 @@ def get_settings() -> dict[str, Any]:
|
||||
"tg_watchdog_heartbeat_every": _parse_int_env("TG_WATCHDOG_HEARTBEAT_EVERY", 30),
|
||||
"tg_disconnect_flap_limit": _parse_int_env("TG_DISCONNECT_FLAP_LIMIT", 3),
|
||||
"tg_disconnect_flap_window": _parse_int_env("TG_DISCONNECT_FLAP_WINDOW", 120),
|
||||
# /ping reports TG as unhealthy once the last successful watchdog probe is older than
|
||||
# this many seconds. Default is derived from the watchdog cadence: it is roughly how
|
||||
# long the watchdog itself would take to give up and trigger a restart —
|
||||
# interval * (failures + 1) + timeout. With the defaults (60,3,10) that is 250s, so a
|
||||
# transient slow probe never flaps /ping, but a genuinely stuck session (no successful
|
||||
# probe for ~4 min) surfaces as 503 before/around the time the watchdog restarts.
|
||||
"tg_ping_unhealthy_after": _parse_int_env(
|
||||
"TG_PING_UNHEALTHY_AFTER",
|
||||
_parse_int_env("TG_WATCHDOG_INTERVAL", 60) * (_parse_int_env("TG_WATCHDOG_FAILURES", 3) + 1)
|
||||
+ _parse_int_env("TG_WATCHDOG_TIMEOUT", 10),
|
||||
),
|
||||
# Media download timeout scales with file size (large videos): the per-download
|
||||
# timeout is clamped to [min, max] seconds, with an effective floor of
|
||||
# `media_download_min_speed` bytes/s (timeout ≈ file_size / min_speed).
|
||||
"media_download_timeout_min": _parse_int_env("MEDIA_DOWNLOAD_TIMEOUT_MIN", 120),
|
||||
"media_download_timeout_max": _parse_int_env("MEDIA_DOWNLOAD_TIMEOUT_MAX", 1800),
|
||||
"media_download_min_speed": _parse_int_env("MEDIA_DOWNLOAD_MIN_SPEED", 256 * 1024),
|
||||
# Size of the asyncio default ThreadPoolExecutor. SQLite/python-magic/pickle/os.walk
|
||||
# all run via asyncio.to_thread; the interpreter default (min(32, cpu+4)) is only 5-6
|
||||
# on a 1-2 CPU container, which starves those under load. 32 gives ample headroom.
|
||||
"io_thread_pool_size": _parse_int_env("IO_THREAD_POOL_SIZE", 32),
|
||||
}
|
||||
|
||||
+9
-3
@@ -28,6 +28,7 @@ services:
|
||||
# MEDIA_DOWNLOAD_TIMEOUT_MIN: 120 # Min per-download timeout, seconds — also the timeout for regular (non-large) files (default: 120)
|
||||
# MEDIA_DOWNLOAD_TIMEOUT_MAX: 1800 # Max per-download timeout, seconds — cap for the largest videos (default: 1800)
|
||||
# MEDIA_DOWNLOAD_MIN_SPEED: 262144 # Assumed floor download speed, bytes/s — large-video timeout ≈ file_size / this, clamped to [MIN,MAX] (default: 262144 = 256 KB/s)
|
||||
# IO_THREAD_POOL_SIZE: 32 # Size of the asyncio default threadpool for blocking I/O (SQLite/python-magic/pickle/os.walk); raise on a busy 1-2 CPU box (default: 32)
|
||||
PYROGRAM_BRIDGE_URL: https://pgbridge.example.com
|
||||
API_PORT: 80
|
||||
TOKEN: ХХХ
|
||||
@@ -52,10 +53,15 @@ services:
|
||||
com.centurylinklabs.watchtower.enable: "true"
|
||||
autoheal: true
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-s", "-f", "-o", "/dev/null", "http://127.0.0.1:80/rss/vvzvlad_lytdybr?limit=1"]
|
||||
interval: 30m
|
||||
# Lightweight process/loop liveness probe. /ping never touches Telegram or the
|
||||
# filesystem, so it answers instantly even while a TG RPC is hung — unlike the old
|
||||
# /rss?limit=1 check, which could exceed the 5s timeout on a cold cache and get the
|
||||
# container restarted mid-download (corrupted temp files). TG liveness is now judged
|
||||
# by the in-process watchdog, which /ping reports via its last-probe age.
|
||||
test: ["CMD", "curl", "-sf", "http://127.0.0.1:80/ping"]
|
||||
interval: 5m
|
||||
timeout: 5s
|
||||
retries: 2
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
start_interval: 5s
|
||||
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
# Стадия 7 — сквозная верификация плана стабилизации
|
||||
|
||||
Дата: 2026-07-05. Ветка: `fix/stage-7-verification` (на базе `fix/stage-6-healthcheck`,
|
||||
кумулятивно содержит стадии 1–6). Эта стадия **ничего не меняет в поведении** прод-кода —
|
||||
только автоматизирует проверки и фиксирует, что должен наблюдать оператор после деплоя.
|
||||
|
||||
## Итог автотестов
|
||||
|
||||
Каноничный прогон (изоляция `config` через `sys.modules`): `python -m pytest tests/`.
|
||||
|
||||
```
|
||||
260 passed
|
||||
```
|
||||
|
||||
Раскладка по стадиям (все зелёные):
|
||||
|
||||
| Стадия | Файл тестов | Кол-во | Что покрывает |
|
||||
|--------|-------------|--------|----------------|
|
||||
| 1 — анти-зависания | `tests/test_stage1_hangs.py` | 6 | таймаут RPC-гейта без утечки пермита; воркер переживает Exception/FloodWait, `task_done` сбалансирован; отмена в spacing-ожидании не теряет пермит |
|
||||
| 2 — статика/большие видео | `tests/test_stage2_static.py` | 15 | атомарный `_download_atomic` (publish-on-rename, чистка `.part.` при таймауте/zero-size/гонке); дедуп конкурентных скачиваний; FloodWait→429; touch mtime у `temp_*`; свипер чистит `.part.`+legacy `.tmp.`; баланс HTTP-семафора |
|
||||
| 3 — FileResponse | `tests/test_stage3_fileresponse.py` | 19 | матрица Range (0-499/500-/-500/за EOF→416/мусор→400/мульти-range→206 multipart); сохранены mtime-touch, delete_after-BackgroundTask, MIME-кэш; чистый ASGI-логгер |
|
||||
| 4 — гигиена event loop | `tests/test_stage4_eventloop.py` | 19 | ленивый `raw_message`; вынос side-effect IO из `process_message` (bulk upsert media id); рендер-пайплайн ушёл в поток без create_task/get_running_loop; XSS вычищен во всех 4 выводах ровно одним проходом |
|
||||
| 5 — батчинг SQLite | `tests/test_stage5_sqlite.py` | 8 | кэш-хит пишет в аккумулятор, а не в SQLite; flush→DB; snapshot-then-clear не теряет поздние апдейты; re-queue при сбое; str(channel)-ключи |
|
||||
| 6 — healthcheck | `tests/test_stage6_healthcheck.py` | 11 | `/ping` 200/503 по connected+age; ноль TG RPC; watchdog отключён→чистая проверка коннекта; `watchdog_last_ok_age()` |
|
||||
| 7 — интеграция | `tests/test_stage7_integration.py` | 8 | сквозные сценарии (см. ниже) |
|
||||
| — регрессия парсера | `tests/test_postparser_*.py` | 174 | заголовки/флаги/автор (существовавшие до плана) |
|
||||
|
||||
## Новые интеграционные тесты стадии 7 (DoD → сквозное доказательство)
|
||||
|
||||
`tests/test_stage7_integration.py` драйвит **реальные точки входа** (`get_media`, `ping`,
|
||||
flush), чтобы поймать регрессии, которые проявляются только при взаимодействии стадий:
|
||||
|
||||
1. **Range на уровне маршрута `/media`** — `test_media_route_range_prefix_0_99`,
|
||||
`_range_suffix_last_100`, `_range_unsatisfiable_416`. Прогоняют `get_media` (кэш-хит) →
|
||||
`prepare_file_response` → `FileResponse` через реальный ASGI: `bytes=0-99`→206 с точным
|
||||
Content-Range/длиной и байтами, `bytes=-100`→206 (суффикс), `bytes=999999999-`→416 (`*/size`).
|
||||
Ловит: если кэш-хит перестанет доходить до FileResponse (ре-буферизация тела, ручные
|
||||
заголовки) или сломается связка digest-гейта — Range перестанет работать. (Стадия 3
|
||||
пинит это на `prepare_file_response` напрямую; здесь — сшивка стадий 2+3.)
|
||||
2. **`/ping` быстр и без RPC при висящей операции** — `test_ping_prompt_and_rpc_free_while_slow_op_pending`
|
||||
и `_reports_degraded_promptly_while_slow_op_pending`. Пока фейковая медленная корутина
|
||||
(модель зависшего hot-path) висит на `Event`, `ping()` возвращает 200/503 корректно и
|
||||
**до** завершения медленной операции, при нуле вызовов `get_me`/`safe_get_messages`.
|
||||
Ловит: рекаплинг `/ping` к TG RPC или к любому awaitable, который может застопориться.
|
||||
3. **Дедуп + очистка при disconnect через реальный `get_media`** —
|
||||
`test_get_media_concurrent_shares_one_download_and_drains_registry` (2 конкурентных
|
||||
запроса → одна скачка, `_inflight` пуст) и `_request_cancel_does_not_stick_registry_or_hang_sibling`
|
||||
(отмена запроса-«клиента» не оставляет застрявший ключ и не вешает соседа). Ловит:
|
||||
возврат скачивания в корутину запроса (отмена убила бы download) или потерю `finally`-pop.
|
||||
4. **str(channel)-ключ access-time end-to-end** — `test_media_cache_hit_flush_updates_str_keyed_db_row`.
|
||||
Кэш-хит `/media` записывает str-ключ, flush обновляет ту же строку в реальной SQLite
|
||||
(hit→flush→DB). Сшивает две половины: ключ, который пишет hot-path, и WHERE, по которому
|
||||
апдейтит flush. Ловит любое расхождение ключа аккумулятора и WHERE bulk-UPDATE (напр.
|
||||
перестановку колонок в SQL — проверено мутацией: тест краснеет, `added` остаётся stale).
|
||||
(int/str-хазард самого канала живёт в `download_media_file` и закрыт стадией 5 —
|
||||
`test_cache_hit_keys_channel_as_str`; здесь покрыта связка get_media+flush.)
|
||||
|
||||
## Соответствие DoD стадий → доказательство
|
||||
|
||||
- **DoD 1** «ни один путь не ждёт TG без таймаута; воркер не умирает молча» →
|
||||
`test_stage1_hangs.py` (гейт-таймаут, живучесть воркера). Живая проверка supervision
|
||||
под нагрузкой — **наблюдение оператора** (лог `supervised_task_crashed/…_exited`).
|
||||
- **DoD 2** «клиенту никогда не отдаётся неполный файл; флуд→429; обрезки не живут >часа» →
|
||||
`test_stage2_static.py` целиком + интеграционный дедуп-тест стадии 7.
|
||||
- **DoD 3** «Range-тесты зелёные; поведение эндпоинта эквивалентно (± RFC-допущения)» →
|
||||
`test_stage3_fileresponse.py` + Range на уровне `/media` (стадия 7). «Нет потока-на-чанк»
|
||||
— **наблюдение оператора** (нагрузочный запрос большого файла, число io-потоков).
|
||||
- **DoD 4** «генерация 100-сообщ. фида не блокирует луп (параллельный /ping <100 мс); XSS
|
||||
зелёный; media id сохраняются» → `test_stage4_eventloop.py` (рендер в потоке, XSS,
|
||||
bulk upsert) + `/ping`-decoupling стадии 7. Живой замер «<100 мс на проде» —
|
||||
**наблюдение оператора**.
|
||||
- **DoD 5** «на кэш-хит /media ноль обращений к SQLite; фоновая запись раз в минуту» →
|
||||
`test_stage5_sqlite.py` + str-ключ end-to-end стадии 7.
|
||||
- **DoD 6** «во время зависшего TG RPC /ping отвечает мгновенно (503); контейнер не
|
||||
рестартится от медленного фида» → `test_stage6_healthcheck.py` + `/ping`-under-slow-op
|
||||
стадии 7. Отсутствие autoheal-рестартов на холодном кэше — **наблюдение оператора**.
|
||||
|
||||
## Ручные curl-сценарии для оператора (пост-деплой, локально в контейнере)
|
||||
|
||||
Подставить реальный `{channel}/{post_id}/{fid}/{digest}` (валидная подпись обязательна).
|
||||
|
||||
```bash
|
||||
# 1. Range-тройка (ожидания: 206 / 206 / 416)
|
||||
curl -s -D- -o /dev/null -H "Range: bytes=0-99" "http://127.0.0.1:80/media/{ch}/{pid}/{fid}/{digest}"
|
||||
curl -s -D- -o /dev/null -H "Range: bytes=-100" "http://127.0.0.1:80/media/{ch}/{pid}/{fid}/{digest}"
|
||||
curl -s -D- -o /dev/null -H "Range: bytes=999999999-" "http://127.0.0.1:80/media/{ch}/{pid}/{fid}/{digest}"
|
||||
|
||||
# 2. Параллельные запросы одного БОЛЬШОГО видео (>100 MB), пока не в кэше:
|
||||
# оба должны получить ПОЛНЫЙ файл (одинаковый размер), без частичной отдачи.
|
||||
URL="http://127.0.0.1:80/media/{ch}/{pid}/{bigfid}/{digest}"
|
||||
curl -s -o /tmp/a "$URL" & curl -s -o /tmp/b "$URL" & wait
|
||||
ls -l /tmp/a /tmp/b # размеры совпадают и равны полному файлу; на диске нет *.part.*
|
||||
|
||||
# 3. Обрыв на середине стрима — нет утечки тасков/фд:
|
||||
# считать fd до/после и убедиться, что не растут монотонно.
|
||||
lsof -p $(pgrep -f api_server) | wc -l # baseline
|
||||
timeout 2 curl -s -o /dev/null "$URL"; sleep 5 # оборвать скачку на середине (несколько раз)
|
||||
lsof -p $(pgrep -f api_server) | wc -l # не выросло относительно baseline
|
||||
|
||||
# 4. Генерация большого фида + параллельный /ping (<100 мс во время генерации):
|
||||
curl -s -o /dev/null "http://127.0.0.1:80/rss/{big_channel}" &
|
||||
for i in $(seq 1 20); do curl -s -o /dev/null -w "%{time_total}\n" "http://127.0.0.1:80/ping"; done
|
||||
wait # все замеры /ping должны быть заметно < 0.100 s
|
||||
```
|
||||
|
||||
> Сценарии 2 и 3 (реальная скачка большого видео + реальный подсчёт fd через `lsof`) в
|
||||
> headless-тестах **намеренно не подделаны** — им нужен живой сервер, реальные загрузки и
|
||||
> реальные файловые дескрипторы. Дедуп-инвариант и disconnect-очистка доказаны на уровне
|
||||
> реестра/`get_media` (тесты стадии 7 №3), но «нет утечки fd на проде» проверяет оператор.
|
||||
|
||||
## Diag-логи для наблюдения после деплоя
|
||||
|
||||
Ожидаемая динамика (сравнить с до-деплойным baseline):
|
||||
|
||||
- `diag_semaphore_wait` — ожидание HTTP-семафора должно **упасть** (реже/меньше секунд).
|
||||
- `diag_download_timing` — время скачивания стабилизируется; нет длинных «зависаний».
|
||||
- `diag_sanitize_slow` — почти **исчезнуть** (один sanitize на выходную границу, стадия 4).
|
||||
- `watchdog: heartbeat` — **продолжаются** штатно (живость TG-сессии).
|
||||
- На момент FloodWait — **нет всплеска 404** на `/media`; вместо этого `media_flood_wait`
|
||||
и ответы **429** с `Retry-After` (стадия 2.3).
|
||||
- Признаки supervision: любые `supervised_task_crashed` / `supervised_task_exited` на
|
||||
CRITICAL — сигнал разобраться, но задача при этом перезапускается (не тихая смерть).
|
||||
|
||||
## План отката (rollback)
|
||||
|
||||
Единица отката — **стадия целиком**, не отдельный коммит. Каждая стадия живёт на своей
|
||||
ветке/PR (`fix/stage-1-hangs` … `fix/stage-7-verification`) и, как правило, состоит из
|
||||
**двух коммитов**: feature-коммит + фикс review-раунда (последний нередко чинит реальный
|
||||
баг — напр. fail-closed XSS в стадии 4, watchdog-gate в стадии 6). Поэтому `git revert`
|
||||
одного коммита осиротит фикс review-раунда и даст несогласованный откат. Откатывать нужно
|
||||
на гранулярности стадии; порядок стадий = порядок деплоя.
|
||||
|
||||
- **Как откатывать стадию** — зависит от того, как ветки влиты в `fix/stability`/`main`:
|
||||
- если стадия влита **squash-merge** (один коммит на стадию) — `git revert <squash-commit>`
|
||||
откатывает её целиком, однозначно;
|
||||
- если стадия влита **merge-коммитом** — `git revert -m 1 <merge-commit>` откатывает весь
|
||||
вклад ветки (оба коммита) разом;
|
||||
- если история линейная (rebase/fast-forward) — реверт **всех** коммитов стадии
|
||||
(`git revert <feature>..<review-fix>` включительно), а не одного.
|
||||
- Стадии 1 и 2 — низкорисковые, деплоятся первыми; откатываются независимо от остальных.
|
||||
- Стадия 3 (FileResponse) независима — реверт возвращает прежний ручной стриминг.
|
||||
- Стадия 4 требует, чтобы 4.2 откатывалась не позже 4.3 (иначе рендер-в-потоке остался бы
|
||||
без безопасного пути сохранения media id) — откатывать стадию 4 целиком.
|
||||
- Стадии 5 и 6 независимы, откатываются по отдельности.
|
||||
- Стадия 7 — только тесты и этот документ: реверт ничего не меняет в проде.
|
||||
|
||||
Прод-деплой, живое наблюдение diag-логов и обновление прод-compose (healthcheck→`/ping`,
|
||||
вне репозитория) — **зона ответственности оператора** (пункты 3–4 плана стадии 7).
|
||||
+39
@@ -69,6 +69,24 @@ def upsert_media_file_id_sync(db_path: str, channel: str, post_id: int, file_uni
|
||||
)
|
||||
|
||||
|
||||
def upsert_media_file_ids_bulk_sync(db_path: str, entries: List[tuple]) -> None:
|
||||
"""Insert or replace multiple media file ID records in a single transaction.
|
||||
|
||||
entries: iterable of (channel, post_id, file_unique_id, added) tuples.
|
||||
Uses executemany for batched upserts (one connection, one commit).
|
||||
"""
|
||||
if not entries:
|
||||
return
|
||||
with _db_connection(db_path) as conn:
|
||||
conn.executemany(
|
||||
"""INSERT INTO media_file_ids (channel, post_id, file_unique_id, added)
|
||||
VALUES (?, ?, ?, ?)
|
||||
ON CONFLICT(channel, post_id, file_unique_id)
|
||||
DO UPDATE SET added = excluded.added""",
|
||||
entries,
|
||||
)
|
||||
|
||||
|
||||
def update_media_file_access_sync(db_path: str, channel: str, post_id: int, file_unique_id: str, added: float) -> None:
|
||||
"""Update the access timestamp for an existing media file ID record."""
|
||||
with _db_connection(db_path) as conn:
|
||||
@@ -78,6 +96,27 @@ def update_media_file_access_sync(db_path: str, channel: str, post_id: int, file
|
||||
)
|
||||
|
||||
|
||||
def update_media_file_access_bulk_sync(db_path: str, entries: List[tuple]) -> None:
|
||||
"""Update the access timestamp for multiple existing media file ID records.
|
||||
|
||||
entries: iterable of (channel, post_id, file_unique_id, added) tuples.
|
||||
Uses executemany (one connection, one commit) so a batch of cache-hit access
|
||||
updates costs a single SQLite transaction instead of one connect+UPDATE per hit.
|
||||
Rows that do not exist are simply not matched by the WHERE clause (no-op), mirroring
|
||||
the single-row update_media_file_access_sync. An empty batch is a no-op.
|
||||
"""
|
||||
if not entries:
|
||||
return
|
||||
with _db_connection(db_path) as conn:
|
||||
conn.executemany(
|
||||
"UPDATE media_file_ids SET added = ? WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
|
||||
# Reorder each (channel, post_id, file_unique_id, added) tuple to match the
|
||||
# UPDATE's placeholder order (added first, then the WHERE key columns).
|
||||
[(added, channel, post_id, file_unique_id)
|
||||
for (channel, post_id, file_unique_id, added) in entries],
|
||||
)
|
||||
|
||||
|
||||
def get_all_media_file_ids_sync(db_path: str) -> List[dict]:
|
||||
"""Return all rows from media_file_ids as a list of dicts."""
|
||||
with _db_connection(db_path) as conn:
|
||||
|
||||
+87
-42
@@ -21,16 +21,13 @@ from pyrogram.enums import MessageMediaType
|
||||
from bleach.css_sanitizer import CSSSanitizer
|
||||
from bleach import clean as HTMLSanitizer
|
||||
from config import get_settings
|
||||
from file_io import upsert_media_file_id_sync, DB_PATH
|
||||
from file_io import upsert_media_file_ids_bulk_sync, DB_PATH
|
||||
from url_signer import generate_media_digest
|
||||
|
||||
Config = get_settings()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Module-level counter for in-flight persist tasks (used only in diagnostics)
|
||||
_persist_pending_count = 0
|
||||
|
||||
|
||||
#tests
|
||||
#http://127.0.0.1:8000/post/html/DragorWW_space/114 — video
|
||||
@@ -65,6 +62,12 @@ _persist_pending_count = 0
|
||||
class PostParser:
|
||||
def __init__(self, client):
|
||||
self.client = client
|
||||
# Media file-id records collected during rendering. _save_media_file_ids
|
||||
# appends (channel, post_id, file_unique_id, ts) tuples here instead of
|
||||
# touching asyncio/DB directly (see 4.2). A fresh PostParser is created per
|
||||
# request and rendering runs in a single thread, so this list is thread-safe.
|
||||
# The caller flushes it once via upsert_media_file_ids_bulk_sync after render.
|
||||
self._pending_media_ids: List[tuple] = []
|
||||
|
||||
@staticmethod
|
||||
def get_all_possible_flags() -> List[str]:
|
||||
@@ -106,7 +109,14 @@ class PostParser:
|
||||
logger.error(f"post_not_found_or_empty: channel {prepared_channel_id}, post_id {post_id}")
|
||||
return None
|
||||
|
||||
processed_message = self.process_message(message)
|
||||
# Single-post outputs need sanitized body/footer (no whole-feed pass exists here).
|
||||
# raw_message is only needed for JSON output and for debug HTML.
|
||||
include_raw = (output_type == 'json') or debug
|
||||
processed_message = self.process_message(message, include_raw=include_raw, sanitize=True)
|
||||
|
||||
# Flush media file-id records collected during rendering with a single bulk upsert.
|
||||
await self._flush_pending_media_ids()
|
||||
|
||||
if output_type == 'html':
|
||||
return self._format_html(processed_message, debug)
|
||||
elif output_type == 'json':
|
||||
@@ -480,11 +490,15 @@ class PostParser:
|
||||
html_content.append(f'<div class="message-body">{data["html"]["body"]}</div>')
|
||||
html_content.append(f'<div class="message-footer">{data["html"]["footer"]}</div>')
|
||||
|
||||
# Add raw JSON debug output if debug is enabled
|
||||
# Add raw JSON debug output if debug is enabled.
|
||||
# raw_message is the full str(message) serialization and may contain user-controlled
|
||||
# text with HTML/JS — escape it before dropping it into the <pre> (bleach can't run
|
||||
# here because <pre> is not in the allowed-tags whitelist and would be stripped).
|
||||
if debug:
|
||||
raw_escaped = html.escape(str(data.get("raw_message", "")))
|
||||
html_content.append(f'<pre class="debug-json" style="background: #f5f5f5;'
|
||||
f'padding: 10px; margin-top: 20px; overflow-x: auto; font-size: 10px;'
|
||||
f'white-space: pre-wrap;">{data["raw_message"]}</pre>')
|
||||
f'white-space: pre-wrap;">{raw_escaped}</pre>')
|
||||
html_data = '\n'.join(html_content)
|
||||
return html_data
|
||||
|
||||
@@ -500,11 +514,39 @@ class PostParser:
|
||||
|
||||
return return_html
|
||||
|
||||
def process_message(self, message: Message) -> Dict[Any, Any]:
|
||||
# Compute html body once — avoids triple _generate_html_body calls
|
||||
def process_message(self, message: Message, include_raw: bool = False, sanitize: bool = True) -> Dict[Any, Any]:
|
||||
"""Build the processed representation of a message.
|
||||
|
||||
Args:
|
||||
message: the Pyrogram message.
|
||||
include_raw: when True, compute the (expensive) full ``str(message)``
|
||||
serialization into ``result['raw_message']``. Only JSON output and
|
||||
debug HTML need it; feed generation must pass False to avoid
|
||||
serializing every post.
|
||||
sanitize: when True, run the html body and footer through a single
|
||||
bleach pass each. Single-post HTML and JSON need this (there is no
|
||||
whole-feed pass on those paths). Feed generation passes False and
|
||||
relies on the final sanitize in rss_generator (per-post for RSS,
|
||||
whole-feed for HTML), so no
|
||||
fragment is sanitized more than once.
|
||||
"""
|
||||
# Compute html body once — avoids triple _generate_html_body calls.
|
||||
# The internal per-fragment sanitize passes were removed (4.4); sanitize
|
||||
# exactly once per output boundary here when requested.
|
||||
html_body = self._generate_html_body(message)
|
||||
# NOTE (stage-4 4.4 consequence): flags are now extracted from the PRE-sanitize
|
||||
# body (the per-fragment sanitize that used to run inside _generate_html_body was
|
||||
# removed). Legitimate links are unaffected — bleach keeps whitelisted
|
||||
# <a href="http(s)://…"> — so link/foreign_channel/mention flags are identical for
|
||||
# normal content. They can differ ONLY for URL-like text that bleach would strip
|
||||
# (e.g. a URL inside a disallowed attribute); flags are non-security (used for
|
||||
# exclude_flags filtering / display), so this edge divergence is accepted rather
|
||||
# than re-adding a per-message sanitize pass that 4.4 deliberately eliminated.
|
||||
flags = self._extract_flags(message, html_body=html_body)
|
||||
footer = self.generate_html_footer(message, flags_list=flags)
|
||||
if sanitize:
|
||||
html_body = self._sanitize_html(html_body)
|
||||
footer = self._sanitize_html(footer)
|
||||
result = {
|
||||
'channel': self.get_channel_username(message),
|
||||
'message_id': message.id,
|
||||
@@ -523,8 +565,10 @@ class PostParser:
|
||||
'media_group_id': message.media_group_id,
|
||||
'service': getattr(message, "service", None),
|
||||
'show_caption_above_media': getattr(message, 'show_caption_above_media', False),
|
||||
'raw_message': str(message)
|
||||
}
|
||||
if include_raw:
|
||||
# Full serialization of the message — expensive; only for JSON/debug.
|
||||
result['raw_message'] = str(message)
|
||||
|
||||
# Add webpage data if available
|
||||
if webpage := getattr(message, "web_page", None):
|
||||
@@ -671,8 +715,10 @@ class PostParser:
|
||||
if message.forward_origin: content_body.append(f"--- Forwarded post end ---") # Forward info end
|
||||
content_body.append(f'</div><br>')
|
||||
|
||||
# NOTE: sanitize is NOT applied here. Sanitization happens exactly once per
|
||||
# output boundary (process_message for single-post/JSON; in rss_generator the
|
||||
# per-post pass for RSS and the whole-feed pass for HTML). See the map (4.4).
|
||||
html_body = '\n'.join(content_body)
|
||||
html_body = self._sanitize_html(html_body)
|
||||
return html_body
|
||||
|
||||
def _generate_html_media(self, message: Message) -> str:
|
||||
@@ -751,8 +797,9 @@ class PostParser:
|
||||
content_media.append(webpage_html)
|
||||
content_media.append('</div>')
|
||||
|
||||
# Not sanitized here — this fragment is embedded in the body and sanitized
|
||||
# once at the output boundary (see the 4.4 sanitize coverage map).
|
||||
html_media = '\n'.join(content_media)
|
||||
html_media = self._sanitize_html(html_media)
|
||||
return html_media
|
||||
|
||||
|
||||
@@ -832,8 +879,9 @@ class PostParser:
|
||||
flags_html = self._format_flags(current_flags)
|
||||
content_footer.append('<br>' + flags_html)
|
||||
|
||||
# Not sanitized here — sanitized once at the output boundary (4.4 coverage map):
|
||||
# process_message for single-post/JSON; per-post (RSS) / whole-feed (HTML) for feeds.
|
||||
html_footer = '\n'.join(content_footer)
|
||||
html_footer = self._sanitize_html(html_footer)
|
||||
return html_footer
|
||||
|
||||
|
||||
@@ -919,8 +967,10 @@ class PostParser:
|
||||
if links:
|
||||
parts.append(' | '.join(links))
|
||||
|
||||
# Raw fragment — embedded in the footer, sanitized once at the output
|
||||
# boundary (see the 4.4 sanitize coverage map). Not sanitized here.
|
||||
result_html = '<br>'.join(parts) if parts else None
|
||||
return self._sanitize_html(result_html) if result_html else None
|
||||
return result_html if result_html else None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"reactions_views_links_error: message_id {message.id}, error {str(e)}")
|
||||
@@ -969,15 +1019,31 @@ class PostParser:
|
||||
logger.error(f"file_id_extraction_error: media_type {message.media}, error {str(e)}")
|
||||
return None
|
||||
|
||||
async def _persist_media_file_id_async(self, channel: str, post_id: int, file_unique_id: str, added: float) -> None:
|
||||
"""Persist a single media file ID record to SQLite (fire-and-forget)."""
|
||||
async def _flush_pending_media_ids(self) -> None:
|
||||
"""Persist media file-id records collected during rendering with ONE bulk upsert.
|
||||
|
||||
Called by the caller (get_post / rss_generator) after rendering completes.
|
||||
Runs the blocking SQLite write in a thread. No-op when nothing was collected.
|
||||
"""
|
||||
entries = self._pending_media_ids
|
||||
if not entries:
|
||||
return
|
||||
try:
|
||||
await asyncio.to_thread(upsert_media_file_id_sync, DB_PATH, channel, post_id, file_unique_id, added)
|
||||
logger.debug(f"persist_media_file_id: upserted {channel}/{post_id}/{file_unique_id}")
|
||||
await asyncio.to_thread(upsert_media_file_ids_bulk_sync, DB_PATH, entries)
|
||||
logger.debug(f"persist_media_file_ids_bulk: upserted {len(entries)} records")
|
||||
except Exception as e:
|
||||
logger.error(f"file_id_save_error: error upserting {channel}/{post_id}/{file_unique_id}, error {str(e)}")
|
||||
logger.error(f"file_id_bulk_save_error: error bulk-upserting {len(entries)} records, error {str(e)}")
|
||||
finally:
|
||||
# Clear regardless of outcome so a retry does not double-persist a stale batch.
|
||||
self._pending_media_ids = []
|
||||
|
||||
def _save_media_file_ids(self, message: Message) -> None:
|
||||
"""Collect a media file-id record for later bulk persistence.
|
||||
|
||||
IMPORTANT (4.2): this runs inside the render thread, so it must NOT touch
|
||||
asyncio or the DB. It only appends to self._pending_media_ids; the caller
|
||||
flushes the batch via _flush_pending_media_ids() after rendering.
|
||||
"""
|
||||
try:
|
||||
channel_username = self.get_channel_username(message)
|
||||
if not channel_username:
|
||||
@@ -1003,29 +1069,8 @@ class PostParser:
|
||||
|
||||
if file_unique_id:
|
||||
added_ts = datetime.now().timestamp()
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
if loop.is_running():
|
||||
global _persist_pending_count
|
||||
task = loop.create_task(
|
||||
self._persist_media_file_id_async(channel_username, message.id, file_unique_id, added_ts)
|
||||
)
|
||||
_persist_pending_count += 1
|
||||
|
||||
def _on_task_done(t):
|
||||
global _persist_pending_count
|
||||
_persist_pending_count -= 1
|
||||
|
||||
task.add_done_callback(_on_task_done)
|
||||
|
||||
if _persist_pending_count > 5:
|
||||
logger.warning(f"persist_task_queue: {_persist_pending_count} pending _persist_media_file_id_async tasks (msg {message.id})")
|
||||
logger.debug(f"persist_task_created: queued for {channel_username}/{message.id}/{file_unique_id}, total pending: {_persist_pending_count}")
|
||||
else:
|
||||
# Fallback sync path (should not occur during normal FastAPI runtime)
|
||||
upsert_media_file_id_sync(DB_PATH, channel_username, message.id, file_unique_id, added_ts)
|
||||
except Exception as e:
|
||||
logger.error(f"file_id_save_error: error saving {channel_username}/{message.id}/{file_unique_id}, error {str(e)}")
|
||||
# Thread-safe: just append; the caller persists the batch.
|
||||
self._pending_media_ids.append((channel_username, message.id, file_unique_id, added_ts))
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"file_id_collection_error: message_id {message.id}, error {str(e)}")
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
fastapi==0.115.8
|
||||
# Pinned explicitly (fastapi only requires a range): the stage-3 FileResponse Range
|
||||
# behaviour and the multipart-byteranges test are tied to this Starlette version.
|
||||
starlette==0.45.3
|
||||
uvicorn==0.34.0
|
||||
python-multipart==0.0.20
|
||||
Kurigram==2.2.22
|
||||
@@ -12,3 +15,4 @@ python-magic==0.4.27
|
||||
bleach[css]==6.1.0
|
||||
types-bleach
|
||||
pytest-asyncio==1.4.0
|
||||
httpx==0.28.1
|
||||
|
||||
+81
-25
@@ -14,6 +14,7 @@ import logging
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
from html import escape as html_escape
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Optional
|
||||
@@ -30,9 +31,12 @@ Config = get_settings()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
async def _create_time_based_media_groups(messages: list[Message], merge_seconds: int = 5) -> list[Message]:
|
||||
def _create_time_based_media_groups(messages: list[Message], merge_seconds: int = 5) -> list[Message]:
|
||||
"""
|
||||
Create media groups based on time difference between messages
|
||||
|
||||
Plain synchronous function (contains no await): runs inside the render thread
|
||||
via _render_pipeline. Must not touch asyncio.
|
||||
"""
|
||||
# Deep-copy the input list to avoid mutating cached Message objects (the cache may
|
||||
# reuse the same objects across calls with different merge_seconds values)
|
||||
@@ -92,9 +96,11 @@ async def _create_time_based_media_groups(messages: list[Message], merge_seconds
|
||||
|
||||
return messages_sorted
|
||||
|
||||
async def _create_messages_groups(messages: list[Message]) -> list[list[Message]]:
|
||||
def _create_messages_groups(messages: list[Message]) -> list[list[Message]]:
|
||||
"""
|
||||
Process messages into formatted posts, handling media groups
|
||||
|
||||
Plain synchronous function (contains no await): runs inside the render thread.
|
||||
"""
|
||||
processing_groups: list[list[Message]] = []
|
||||
media_groups: dict[str | int, list[Message]] = {}
|
||||
@@ -136,9 +142,11 @@ async def _create_messages_groups(messages: list[Message]) -> list[list[Message]
|
||||
|
||||
return processing_groups
|
||||
|
||||
async def _trim_messages_groups(messages_groups: list[list[Message]], limit: int):
|
||||
def _trim_messages_groups(messages_groups: list[list[Message]], limit: int):
|
||||
"""
|
||||
Trim messages groups to limit
|
||||
|
||||
Plain synchronous function (contains no await): runs inside the render thread.
|
||||
"""
|
||||
if len(messages_groups) > limit: # Trim groups if they exceed the specified limit
|
||||
messages_groups = messages_groups[:limit]
|
||||
@@ -193,12 +201,13 @@ def processed_message_to_tg_message(processed_message: dict) -> Message:
|
||||
return tg_message_mock # type: ignore
|
||||
|
||||
|
||||
async def _render_messages_groups(messages_groups: list[list[Message]],
|
||||
post_parser: PostParser,
|
||||
exclude_flags: str | None = None,
|
||||
def _render_messages_groups(messages_groups: list[list[Message]],
|
||||
post_parser: PostParser,
|
||||
exclude_flags: str | None = None,
|
||||
exclude_text: str | None = None):
|
||||
"""
|
||||
Render message groups into HTML format
|
||||
Plain synchronous function (contains no await): runs inside the render thread.
|
||||
Args:
|
||||
messages_groups: List of message groups (each group is a list of messages)
|
||||
post_parser: PostParser instance
|
||||
@@ -213,7 +222,9 @@ async def _render_messages_groups(messages_groups: list[list[Message]],
|
||||
try:
|
||||
if len(group) == 1: # Single message - simple case
|
||||
one_message = group[0]
|
||||
message_data = post_parser.process_message(one_message)
|
||||
# Feed path: raw_message not needed and sanitize deferred to the final
|
||||
# whole-feed pass, so each fragment is sanitized exactly once.
|
||||
message_data = post_parser.process_message(one_message, include_raw=False, sanitize=False)
|
||||
html_parts = [
|
||||
f'<div class="message-body">{message_data["html"]["body"]}</div>',
|
||||
f'<div class="message-footer">{message_data["html"]["footer"]}</div>'
|
||||
@@ -228,7 +239,7 @@ async def _render_messages_groups(messages_groups: list[list[Message]],
|
||||
'flags': message_data['flags']
|
||||
})
|
||||
else: # Multiple messages in group - merge text and html body
|
||||
processed_messages = [post_parser.process_message(msg) for msg in group]
|
||||
processed_messages = [post_parser.process_message(msg, include_raw=False, sanitize=False) for msg in group]
|
||||
|
||||
# Determine main message for header/footer/title
|
||||
main_message = next(
|
||||
@@ -309,7 +320,31 @@ async def _render_messages_groups(messages_groups: list[list[Message]],
|
||||
rendered_posts.sort(key=lambda x: x['date'] if x['date'] is not None else 0.0, reverse=True)
|
||||
return rendered_posts
|
||||
|
||||
async def generate_channel_rss(channel: str | int,
|
||||
|
||||
def _render_pipeline(messages: list[Message],
|
||||
post_parser: PostParser,
|
||||
limit: int,
|
||||
exclude_flags: str | None,
|
||||
exclude_text: str | None,
|
||||
merge_seconds: int,
|
||||
time_based_merge: bool):
|
||||
"""
|
||||
Full synchronous feed render pipeline (grouping + trimming + rendering).
|
||||
|
||||
Runs entirely in a worker thread via a single asyncio.to_thread call. It contains
|
||||
NO await, NO asyncio primitives, NO create_task/get_running_loop — all the CPU-heavy
|
||||
work (deepcopy, grouping, bleach-free rendering) happens here off the event loop.
|
||||
Media file-id records are accumulated on post_parser._pending_media_ids and flushed
|
||||
by the caller after this returns.
|
||||
"""
|
||||
if time_based_merge:
|
||||
messages = _create_time_based_media_groups(messages, merge_seconds)
|
||||
message_groups = _create_messages_groups(messages)
|
||||
message_groups = _trim_messages_groups(message_groups, limit)
|
||||
return _render_messages_groups(message_groups, post_parser, exclude_flags, exclude_text)
|
||||
|
||||
|
||||
async def generate_channel_rss(channel: str | int,
|
||||
client: Client,
|
||||
limit: int = 20,
|
||||
exclude_flags: str | None = None,
|
||||
@@ -389,14 +424,23 @@ async def generate_channel_rss(channel: str | int,
|
||||
messages_elapsed = time.time() - messages_start_time
|
||||
logger.debug(f"rss_messages_retrieval_timing: channel {channel}, {len(messages)} messages retrieved in {messages_elapsed:.3f} seconds")
|
||||
|
||||
# Process messages into groups and render them
|
||||
# Process messages into groups and render them.
|
||||
# The whole grouping/trimming/rendering pipeline is CPU-heavy (deepcopy +
|
||||
# per-message rendering) and contains no await, so run it in ONE worker
|
||||
# thread to keep the event loop responsive.
|
||||
processing_start_time = time.time()
|
||||
if Config['time_based_merge']:
|
||||
messages = await _create_time_based_media_groups(messages, merge_seconds)
|
||||
message_groups = await _create_messages_groups(messages)
|
||||
message_groups = await _trim_messages_groups(message_groups, limit)
|
||||
final_posts = await _render_messages_groups(message_groups, post_parser, exclude_flags, exclude_text)
|
||||
|
||||
try:
|
||||
final_posts = await asyncio.to_thread(
|
||||
_render_pipeline, messages, post_parser, limit,
|
||||
exclude_flags, exclude_text, merge_seconds, Config['time_based_merge'],
|
||||
)
|
||||
finally:
|
||||
# Persist media file-ids collected during rendering with a single bulk
|
||||
# upsert — in a finally so a partial render still records what it collected
|
||||
# (the flush is best-effort and swallows its own errors, so it cannot mask a
|
||||
# render exception).
|
||||
await post_parser._flush_pending_media_ids()
|
||||
|
||||
processing_elapsed = time.time() - processing_start_time
|
||||
logger.debug(f"rss_messages_processing_timing: channel {channel}, {len(final_posts)} posts processed in {processing_elapsed:.3f} seconds")
|
||||
|
||||
@@ -443,8 +487,13 @@ async def generate_channel_rss(channel: str | int,
|
||||
)
|
||||
sanitized_html = await asyncio.to_thread(_sanitize_sync, post['html'])
|
||||
except Exception as e:
|
||||
# FAIL CLOSED: since 4.4 removed the per-fragment passes, post['html'] is
|
||||
# raw channel-controlled HTML, and this is its ONLY sanitize. If bleach
|
||||
# itself throws (e.g. RecursionError on pathological nested HTML), do NOT
|
||||
# emit the raw payload (that would be stored XSS) — html.escape it so the
|
||||
# content survives as inert text.
|
||||
logger.error(f"rss_html_sanitization_error: channel {channel}, message_id {post['message_id']}, error {str(e)}")
|
||||
sanitized_html = post['html']
|
||||
sanitized_html = html_escape(post['html'])
|
||||
fe.content(content=sanitized_html, type='CDATA')
|
||||
|
||||
if post['date'] is not None:
|
||||
@@ -601,16 +650,19 @@ async def generate_channel_html(channel: str | int,
|
||||
enrichment_elapsed = time.time() - enrichment_start_time
|
||||
logger.debug(f"html_reply_enrichment_timing: channel {channel}, replies enriched in {enrichment_elapsed:.3f} seconds")
|
||||
|
||||
# Process messages into groups and render them
|
||||
# Process messages into groups and render them in ONE worker thread
|
||||
# (CPU-heavy, no await) to keep the event loop responsive.
|
||||
processing_start_time = time.time()
|
||||
if Config['time_based_merge']:
|
||||
messages = await _create_time_based_media_groups(messages, merge_seconds)
|
||||
try:
|
||||
final_posts = await asyncio.to_thread(
|
||||
_render_pipeline, messages, post_parser, limit,
|
||||
exclude_flags, exclude_text, merge_seconds, Config['time_based_merge'],
|
||||
)
|
||||
finally:
|
||||
# Persist media file-ids collected during rendering with a single bulk
|
||||
# upsert — in a finally so a partial render still records what it collected.
|
||||
await post_parser._flush_pending_media_ids()
|
||||
|
||||
# Process messages into groups and render them
|
||||
message_groups = await _create_messages_groups(messages)
|
||||
message_groups = await _trim_messages_groups(message_groups, limit)
|
||||
final_posts = await _render_messages_groups(message_groups, post_parser, exclude_flags, exclude_text)
|
||||
|
||||
processing_elapsed = time.time() - processing_start_time
|
||||
logger.debug(f"html_messages_processing_timing: channel {channel}, {len(final_posts)} posts processed in {processing_elapsed:.3f} seconds")
|
||||
|
||||
@@ -646,7 +698,11 @@ async def generate_channel_html(channel: str | int,
|
||||
)
|
||||
html = await asyncio.to_thread(_sanitize_sync, html)
|
||||
except Exception as e:
|
||||
# FAIL CLOSED (see the RSS path): `html` is now the raw concatenated feed
|
||||
# (4.4 removed the per-fragment passes). If bleach throws, html.escape the
|
||||
# whole feed rather than returning raw channel HTML/JS to the client.
|
||||
logger.error(f"html_final_sanitization_error: channel {channel}, error {str(e)}")
|
||||
html = html_escape(html)
|
||||
|
||||
html_gen_elapsed = time.time() - html_gen_start_time
|
||||
logger.debug(f"html_generation_timing: channel {channel}, HTML generated in {html_gen_elapsed:.3f} seconds")
|
||||
|
||||
@@ -259,6 +259,17 @@ class TelegramClient:
|
||||
# Emergency termination
|
||||
os._exit(1)
|
||||
|
||||
def watchdog_last_ok_age(self) -> float | None:
|
||||
"""Seconds since the last successful watchdog probe, or None if none succeeded yet.
|
||||
|
||||
Reads only the already-recorded monotonic timestamp set by the watchdog loop; it
|
||||
never issues a Telegram RPC, so it is safe to call from the hot /ping path even
|
||||
while a real RPC is hung.
|
||||
"""
|
||||
if self._wd_last_ok_monotonic is None:
|
||||
return None
|
||||
return time.monotonic() - self._wd_last_ok_monotonic
|
||||
|
||||
async def safe_get_messages(self, channel_id, post_id, max_retries=2):
|
||||
"""Wrapper with retry logic for auth errors"""
|
||||
for attempt in range(max_retries):
|
||||
|
||||
@@ -33,7 +33,9 @@ def get_settings():
|
||||
"tg_watchdog_heartbeat_every": 30,
|
||||
"tg_disconnect_flap_limit": 3,
|
||||
"tg_disconnect_flap_window": 120,
|
||||
"tg_ping_unhealthy_after": 250,
|
||||
"media_download_timeout_min": 120,
|
||||
"media_download_timeout_max": 1800,
|
||||
"media_download_min_speed": 256 * 1024,
|
||||
"io_thread_pool_size": 32,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,383 @@
|
||||
# flake8: noqa
|
||||
# pylint: disable=protected-access, missing-function-docstring, missing-class-docstring
|
||||
# pylint: disable=redefined-outer-name, logging-fstring-interpolation, line-too-long
|
||||
# pylance: disable=reportMissingImports, reportMissingModuleSource
|
||||
"""
|
||||
Stage 3 (FileResponse + HTTP-layer cleanup) regression tests.
|
||||
|
||||
Covers:
|
||||
- The HTTP Range matrix asserted against the FINAL (Starlette FileResponse) behavior,
|
||||
each case documented, including the RFC-7233-permitted deltas vs the old hand-rolled
|
||||
streaming (garbage header -> 400 not 416; multi-range -> proper multipart 206 not 416;
|
||||
416 Content-Range now `*/size`; ETag/Last-Modified now present; ASCII filename no longer
|
||||
gets a redundant filename*= form).
|
||||
- Stage-2 behavior preserved through the rewrite: a served temp_* file still has its mtime
|
||||
refreshed; delete_after still schedules the temp-file BackgroundTask (and it actually
|
||||
runs); the media_key MIME cache is still consulted and populated.
|
||||
- The pure-ASGI RequestLoggingMiddleware: a normal request still returns (body intact, not
|
||||
buffered/truncated) and is logged.
|
||||
|
||||
Before/after Range matrix (2048-byte file), old = hand-rolled streaming, new = FileResponse:
|
||||
|
||||
request | old | new (FileResponse)
|
||||
-----------------------+-----------------------------+-----------------------------------
|
||||
no Range | 200 full | 200 full (+ETag/Last-Modified)
|
||||
bytes=0-499 | 206 bytes 0-499/2048 | 206 bytes 0-499/2048 (same bytes)
|
||||
bytes=500- | 206 bytes 500-2047/2048 | 206 bytes 500-2047/2048 (same)
|
||||
bytes=-500 | 206 bytes 1548-2047/2048 | 206 bytes 1548-2047/2048 (same)
|
||||
bytes=999999- (>EOF) | 416 `bytes */2048` | 416 `*/2048` (Starlette omits the
|
||||
| | `bytes ` prefix; RFC-acceptable)
|
||||
garbage header | 416 | 400 Bad Request (RFC 7233 lets a
|
||||
| | server reject/ignore a malformed
|
||||
| | Range; Starlette returns 400)
|
||||
bytes=0-9,20-29 | 416 (old parser bug) | 206 multipart/byteranges (correct)
|
||||
|
||||
All 206 responses that carry data return byte-for-byte identical slices old vs new, so no
|
||||
real regression is hidden behind an "accepted difference".
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
# Add project root to sys.path and mock the config module (same pattern as the other tests).
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings'])
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import api_server
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Helpers
|
||||
# --------------------------------------------------------------------------- #
|
||||
BODY = bytes(range(256)) * 8 # 2048 deterministic bytes
|
||||
SIZE = len(BODY)
|
||||
|
||||
|
||||
def _make_client(file_path, delete_after=False, media_key=None):
|
||||
"""Mount prepare_file_response on a tiny app so it is driven through real ASGI
|
||||
(FileResponse computes Range/206/416 at send time, so it must be exercised via a client)."""
|
||||
app = FastAPI()
|
||||
|
||||
@app.get("/f")
|
||||
async def _serve(request: Request):
|
||||
return await api_server.prepare_file_response(
|
||||
file_path, request=request, delete_after=delete_after, media_key=media_key
|
||||
)
|
||||
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_file(tmp_path):
|
||||
fp = tmp_path / "myfile.bin"
|
||||
fp.write_bytes(BODY)
|
||||
return str(fp)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 3.1 — Range matrix against the FINAL FileResponse behavior.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_no_range_returns_full_200(sample_file):
|
||||
c = _make_client(sample_file)
|
||||
r = c.get("/f")
|
||||
assert r.status_code == 200
|
||||
assert r.content == BODY
|
||||
assert r.headers["content-length"] == str(SIZE)
|
||||
assert r.headers["accept-ranges"] == "bytes"
|
||||
assert r.headers["cache-control"] == "public, max-age=86400, immutable"
|
||||
# Content-Disposition: inline, filename set by FileResponse itself (no manual double-set).
|
||||
assert r.headers["content-disposition"].startswith("inline")
|
||||
assert 'filename="myfile.bin"' in r.headers["content-disposition"]
|
||||
# NEW vs old: FileResponse adds validators. Old streaming had neither.
|
||||
assert r.headers.get("etag")
|
||||
assert r.headers.get("last-modified")
|
||||
|
||||
|
||||
def test_range_prefix_0_499(sample_file):
|
||||
c = _make_client(sample_file)
|
||||
r = c.get("/f", headers={"Range": "bytes=0-499"})
|
||||
assert r.status_code == 206
|
||||
assert r.headers["content-range"] == f"bytes 0-499/{SIZE}"
|
||||
assert r.headers["content-length"] == "500"
|
||||
assert r.content == BODY[:500] # exact slice, identical to old behavior
|
||||
|
||||
|
||||
def test_range_open_ended_500_to_eof(sample_file):
|
||||
c = _make_client(sample_file)
|
||||
r = c.get("/f", headers={"Range": "bytes=500-"})
|
||||
assert r.status_code == 206
|
||||
assert r.headers["content-range"] == f"bytes 500-{SIZE - 1}/{SIZE}"
|
||||
assert r.headers["content-length"] == str(SIZE - 500)
|
||||
assert r.content == BODY[500:]
|
||||
|
||||
|
||||
def test_range_suffix_last_500(sample_file):
|
||||
c = _make_client(sample_file)
|
||||
r = c.get("/f", headers={"Range": "bytes=-500"})
|
||||
assert r.status_code == 206
|
||||
assert r.headers["content-range"] == f"bytes {SIZE - 500}-{SIZE - 1}/{SIZE}"
|
||||
assert r.headers["content-length"] == "500"
|
||||
assert r.content == BODY[-500:]
|
||||
|
||||
|
||||
def test_range_start_past_eof_returns_416(sample_file):
|
||||
c = _make_client(sample_file)
|
||||
r = c.get("/f", headers={"Range": "bytes=999999-"})
|
||||
assert r.status_code == 416
|
||||
# DELTA (RFC-acceptable): Starlette's 416 Content-Range is `*/size` (it omits the
|
||||
# `bytes ` unit prefix the old code emitted). Same information, permitted by RFC 7233.
|
||||
assert r.headers["content-range"] == f"*/{SIZE}"
|
||||
|
||||
|
||||
def test_garbage_range_header_returns_400(sample_file):
|
||||
c = _make_client(sample_file)
|
||||
r = c.get("/f", headers={"Range": "somethinggarbage"})
|
||||
# DELTA (RFC-acceptable): the old hand-rolled parser returned 416 on an unparseable
|
||||
# header; Starlette rejects a malformed Range with 400 Bad Request. RFC 7233 permits a
|
||||
# server to reject/ignore a bad Range; this is a conscious accepted difference, not a
|
||||
# data regression (no bytes are served either way).
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
def test_multi_range_returns_multipart_206(sample_file):
|
||||
c = _make_client(sample_file)
|
||||
r = c.get("/f", headers={"Range": "bytes=0-9,20-29"})
|
||||
# DELTA (improvement): the old parser mis-split multi-ranges and returned 416.
|
||||
# FileResponse serves a proper multipart/byteranges 206 containing both slices.
|
||||
# Starlette 0.45.3 quirk: it advertises the multipart boundary via the `content-range`
|
||||
# header (rather than content-type, which stays the file's own media_type); the body is
|
||||
# a real multipart/byteranges document. We assert on the boundary marker + both slices.
|
||||
assert r.status_code == 206
|
||||
assert r.headers["content-range"].startswith("multipart/byteranges")
|
||||
body = r.content
|
||||
assert BODY[0:10] in body
|
||||
assert BODY[20:30] in body
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 3.1 (stage-2 preserved) — temp_* mtime touch survives the FileResponse swap.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_temp_file_mtime_refreshed_when_stale(tmp_path):
|
||||
# A temp_* file whose mtime is OLDER than the refresh interval gets touched, so the
|
||||
# 1h sweeper won't delete an actively-viewed video.
|
||||
temp_file = tmp_path / "temp_bigvid"
|
||||
temp_file.write_bytes(b"videodata")
|
||||
stale = time.time() - 10000 # >> TEMP_MTIME_REFRESH_INTERVAL (300s)
|
||||
os.utime(temp_file, (stale, stale))
|
||||
|
||||
c = _make_client(str(temp_file))
|
||||
r = c.get("/f")
|
||||
assert r.status_code == 200
|
||||
assert r.content == b"videodata"
|
||||
assert os.path.getmtime(temp_file) > time.time() - 100 # refreshed by the serve
|
||||
|
||||
|
||||
def test_temp_file_mtime_stable_when_fresh(tmp_path):
|
||||
# DEBOUNCE: a temp_* file touched RECENTLY (within the refresh interval) is NOT
|
||||
# re-touched — so FileResponse's mtime-derived ETag stays stable across the rapid
|
||||
# requests of a single resume/seek session and `If-Range` resume keeps working.
|
||||
temp_file = tmp_path / "temp_freshvid"
|
||||
temp_file.write_bytes(b"videodata")
|
||||
recent = time.time() - 30 # << TEMP_MTIME_REFRESH_INTERVAL (300s)
|
||||
os.utime(temp_file, (recent, recent))
|
||||
mtime_before = os.path.getmtime(temp_file)
|
||||
|
||||
c = _make_client(str(temp_file))
|
||||
r1 = c.get("/f")
|
||||
r2 = c.get("/f")
|
||||
assert r1.status_code == 200 and r2.status_code == 200
|
||||
# mtime unchanged -> the ETag is identical across the two serves.
|
||||
assert os.path.getmtime(temp_file) == mtime_before
|
||||
assert r1.headers.get("etag") == r2.headers.get("etag")
|
||||
|
||||
|
||||
def test_non_temp_file_mtime_untouched(tmp_path):
|
||||
normal = tmp_path / "regularfile"
|
||||
normal.write_bytes(b"data")
|
||||
stale = time.time() - 10000
|
||||
os.utime(normal, (stale, stale))
|
||||
|
||||
c = _make_client(str(normal))
|
||||
c.get("/f")
|
||||
assert os.path.getmtime(normal) < time.time() - 5000 # left alone
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 3.1 (stage-2 preserved) — delete_after still schedules the temp-file BackgroundTask.
|
||||
# --------------------------------------------------------------------------- #
|
||||
async def test_delete_after_attaches_background_task(tmp_path):
|
||||
fp = tmp_path / "temp_todelete"
|
||||
fp.write_bytes(b"x")
|
||||
|
||||
class _Req:
|
||||
headers = {}
|
||||
|
||||
resp = await api_server.prepare_file_response(str(fp), request=_Req(), delete_after=True)
|
||||
# FileResponse must carry a non-None background that deletes exactly this file.
|
||||
assert resp.background is not None
|
||||
assert resp.background.func is api_server.delayed_delete_file
|
||||
assert resp.background.args == (str(fp),)
|
||||
|
||||
|
||||
async def test_no_delete_after_has_no_background(tmp_path):
|
||||
fp = tmp_path / "keepme"
|
||||
fp.write_bytes(b"x")
|
||||
|
||||
class _Req:
|
||||
headers = {}
|
||||
|
||||
resp = await api_server.prepare_file_response(str(fp), request=_Req(), delete_after=False)
|
||||
assert resp.background is None
|
||||
|
||||
|
||||
def test_delete_after_background_runs_and_removes_file(tmp_path, monkeypatch):
|
||||
fp = tmp_path / "temp_gone"
|
||||
fp.write_bytes(BODY)
|
||||
|
||||
# Patch the module-level deleter to remove immediately (real one sleeps 300s); the
|
||||
# BackgroundTask picks up this reference at prepare_file_response call time.
|
||||
async def _delete_now(path, delay=300):
|
||||
os.remove(path)
|
||||
|
||||
monkeypatch.setattr(api_server, "delayed_delete_file", _delete_now)
|
||||
|
||||
c = _make_client(str(fp), delete_after=True)
|
||||
r = c.get("/f")
|
||||
assert r.status_code == 200
|
||||
assert r.content == BODY
|
||||
# TestClient blocks until the background task has run.
|
||||
assert not os.path.exists(fp)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 3.1 (stage-2 preserved) — media_key MIME cache is consulted and populated.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_media_key_mime_cache_hit_used(sample_file, monkeypatch):
|
||||
calls = {"get": 0, "set": 0, "magic": 0}
|
||||
|
||||
def fake_get(db, ch, pid, fid):
|
||||
calls["get"] += 1
|
||||
return "video/mp4" # cache HIT
|
||||
|
||||
def fake_set(*a, **k):
|
||||
calls["set"] += 1
|
||||
|
||||
def fake_magic(_path):
|
||||
calls["magic"] += 1
|
||||
return "application/octet-stream"
|
||||
|
||||
monkeypatch.setattr(api_server, "get_mime_type_sync", fake_get)
|
||||
monkeypatch.setattr(api_server, "set_mime_type_sync", fake_set)
|
||||
monkeypatch.setattr(api_server.magic_mime, "from_file", fake_magic)
|
||||
|
||||
c = _make_client(sample_file, media_key=("chan", 42, "fid"))
|
||||
r = c.get("/f")
|
||||
assert r.status_code == 200
|
||||
# Cache hit: python-magic never invoked, nothing re-written to the cache, MIME applied.
|
||||
assert calls["get"] == 1
|
||||
assert calls["magic"] == 0
|
||||
assert calls["set"] == 0
|
||||
assert r.headers["content-type"].startswith("video/mp4")
|
||||
|
||||
|
||||
def test_media_key_mime_cache_miss_populates(sample_file, monkeypatch):
|
||||
written = {}
|
||||
|
||||
def fake_get(db, ch, pid, fid):
|
||||
return None # cache MISS
|
||||
|
||||
def fake_set(db, ch, pid, fid, mime):
|
||||
written["mime"] = mime
|
||||
|
||||
def fake_magic(_path):
|
||||
return "image/png"
|
||||
|
||||
monkeypatch.setattr(api_server, "get_mime_type_sync", fake_get)
|
||||
monkeypatch.setattr(api_server, "set_mime_type_sync", fake_set)
|
||||
monkeypatch.setattr(api_server.magic_mime, "from_file", fake_magic)
|
||||
|
||||
c = _make_client(sample_file, media_key=("chan", 42, "fid"))
|
||||
r = c.get("/f")
|
||||
assert r.status_code == 200
|
||||
# Miss -> python-magic result is both applied and persisted for next time.
|
||||
assert written.get("mime") == "image/png"
|
||||
assert r.headers["content-type"].startswith("image/png")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 3.1 — 404 pre-check kept.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_missing_file_returns_404(tmp_path):
|
||||
c = _make_client(str(tmp_path / "does_not_exist"))
|
||||
r = c.get("/f")
|
||||
assert r.status_code == 404
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 3.2 — pure-ASGI RequestLoggingMiddleware: request returns intact and is logged.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_asgi_logging_middleware_passes_body_and_logs(tmp_path, caplog):
|
||||
fp = tmp_path / "streamed.bin"
|
||||
fp.write_bytes(BODY)
|
||||
|
||||
app = FastAPI()
|
||||
app.add_middleware(api_server.RequestLoggingMiddleware)
|
||||
|
||||
@app.get("/f")
|
||||
async def _serve(request: Request):
|
||||
return await api_server.prepare_file_response(str(fp), request=request)
|
||||
|
||||
client = TestClient(app)
|
||||
with caplog.at_level(logging.DEBUG, logger="api_server"):
|
||||
r = client.get("/f")
|
||||
|
||||
# Body flows straight through the middleware — not buffered/truncated.
|
||||
assert r.status_code == 200
|
||||
assert r.content == BODY
|
||||
messages = " ".join(rec.getMessage() for rec in caplog.records)
|
||||
assert "Request: GET /f" in messages
|
||||
assert "Response status: 200" in messages
|
||||
|
||||
|
||||
def test_asgi_logging_middleware_range_still_works(tmp_path):
|
||||
fp = tmp_path / "streamed.bin"
|
||||
fp.write_bytes(BODY)
|
||||
|
||||
app = FastAPI()
|
||||
app.add_middleware(api_server.RequestLoggingMiddleware)
|
||||
|
||||
@app.get("/f")
|
||||
async def _serve(request: Request):
|
||||
return await api_server.prepare_file_response(str(fp), request=request)
|
||||
|
||||
client = TestClient(app)
|
||||
r = client.get("/f", headers={"Range": "bytes=0-99"})
|
||||
# 206 streaming still works through the pure-ASGI middleware (send not buffered).
|
||||
assert r.status_code == 206
|
||||
assert r.content == BODY[:100]
|
||||
|
||||
|
||||
async def test_asgi_logging_middleware_passes_non_http_scope_through(caplog):
|
||||
# The `scope["type"] != "http"` branch (lifespan/websocket) only runs in a real
|
||||
# deploy, never through TestClient's plain request path — so pin it directly:
|
||||
# a non-http scope must delegate to the inner app untouched and log nothing.
|
||||
called = {}
|
||||
|
||||
async def inner(scope, receive, send):
|
||||
called["scope_type"] = scope["type"]
|
||||
|
||||
mw = api_server.RequestLoggingMiddleware(inner)
|
||||
with caplog.at_level(logging.DEBUG, logger="api_server"):
|
||||
await mw({"type": "lifespan"}, None, None)
|
||||
|
||||
assert called["scope_type"] == "lifespan" # delegated to the inner app
|
||||
# Nothing request/response-ish was logged for a non-http scope.
|
||||
messages = " ".join(rec.getMessage() for rec in caplog.records)
|
||||
assert "Request:" not in messages
|
||||
assert "Response status:" not in messages
|
||||
@@ -0,0 +1,471 @@
|
||||
# flake8: noqa
|
||||
# pylint: disable=protected-access, missing-function-docstring, missing-class-docstring
|
||||
# pylint: disable=redefined-outer-name, logging-fstring-interpolation, line-too-long
|
||||
# pylance: disable=reportMissingImports, reportMissingModuleSource
|
||||
"""
|
||||
Stage 4 (event-loop hygiene for feed generation) tests.
|
||||
|
||||
Covers:
|
||||
- 4.1 raw_message laziness: feeds do NOT compute str(message); JSON/debug HTML do.
|
||||
- 4.2 side-effect IO removed from process_message: _save_media_file_ids only appends to
|
||||
self._pending_media_ids; the caller flushes once via upsert_media_file_ids_bulk_sync.
|
||||
Also: the render path contains NO create_task / get_running_loop / to_thread.
|
||||
- 4.3 render pipeline moved into a single thread: the four render functions are now plain
|
||||
sync functions and actually execute off the main thread; deepcopy of a pickled Message
|
||||
does not crash; a 100-message feed generates correctly.
|
||||
- 4.4 sanitize coverage (XSS): a <script> / onerror= / javascript: payload is stripped in
|
||||
ALL outputs — rss, html-feed, single-post html, and json — each with exactly one pass.
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import copy
|
||||
import pickle
|
||||
import asyncio
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings'])
|
||||
|
||||
from pyrogram.enums import MessageMediaType
|
||||
|
||||
import post_parser as pp_module
|
||||
import rss_generator as rss_module
|
||||
from post_parser import PostParser
|
||||
from rss_generator import (
|
||||
generate_channel_rss,
|
||||
generate_channel_html,
|
||||
_render_pipeline,
|
||||
_create_time_based_media_groups,
|
||||
_create_messages_groups,
|
||||
_trim_messages_groups,
|
||||
_render_messages_groups,
|
||||
)
|
||||
|
||||
XSS_PAYLOAD = "<script>alert('xss')</script><img src=x onerror=\"alert(1)\"><a href=\"javascript:alert(2)\">click</a>"
|
||||
|
||||
|
||||
class _Str(str):
|
||||
"""Minimal stand-in for Pyrogram's Str: .html returns the raw string unchanged,
|
||||
so a malicious payload reaches the pre-sanitization html body just like real text
|
||||
carrying entities would."""
|
||||
@property
|
||||
def html(self):
|
||||
return str(self)
|
||||
|
||||
|
||||
def make_message(mid, text=None, media=None, photo_uid=None, username="testchan",
|
||||
date=None):
|
||||
m = SimpleNamespace()
|
||||
m.id = mid
|
||||
m.date = date or datetime(2024, 1, 1, 12, 0, mid % 60, tzinfo=timezone.utc)
|
||||
m.text = _Str(text) if text is not None else None
|
||||
m.caption = None
|
||||
m.media = media
|
||||
m.web_page = None
|
||||
m.poll = None
|
||||
m.service = None
|
||||
m.forward_origin = None
|
||||
m.reply_to_message = None
|
||||
m.reply_to_message_id = None
|
||||
m.sender_chat = None
|
||||
m.from_user = None
|
||||
m.reactions = None
|
||||
m.views = 100
|
||||
m.media_group_id = None
|
||||
m.show_caption_above_media = False
|
||||
m.chat = SimpleNamespace(id=-1001234567890, username=username)
|
||||
# media sub-objects default to None
|
||||
for attr in ("photo", "video", "document", "audio", "voice",
|
||||
"video_note", "animation", "sticker"):
|
||||
setattr(m, attr, None)
|
||||
if media == MessageMediaType.PHOTO and photo_uid:
|
||||
m.photo = SimpleNamespace(file_unique_id=photo_uid)
|
||||
return m
|
||||
|
||||
|
||||
def _co_names(func):
|
||||
return set(func.__code__.co_names)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4.3 — render functions are plain sync and run off the main thread
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_render_functions_are_sync():
|
||||
for fn in (_create_time_based_media_groups, _create_messages_groups,
|
||||
_trim_messages_groups, _render_messages_groups, _render_pipeline):
|
||||
assert not asyncio.iscoroutinefunction(fn), f"{fn.__name__} must be a plain sync function"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_runs_in_worker_thread(monkeypatch):
|
||||
main_ident = threading.get_ident()
|
||||
seen = {}
|
||||
|
||||
real_render = rss_module._render_messages_groups
|
||||
|
||||
def spy(*args, **kwargs):
|
||||
seen["ident"] = threading.get_ident()
|
||||
return real_render(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(rss_module, "_render_messages_groups", spy)
|
||||
|
||||
async def fake_get_chat(client, channel):
|
||||
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
|
||||
|
||||
async def fake_get_history(client, channel, limit=20):
|
||||
return [make_message(i, text=f"post {i}") for i in range(5)]
|
||||
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
|
||||
|
||||
await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
|
||||
assert "ident" in seen
|
||||
assert seen["ident"] != main_ident, "render pipeline must run in a worker thread, not the loop thread"
|
||||
|
||||
|
||||
def test_deepcopy_of_pickled_message_does_not_crash():
|
||||
from pyrogram.types import Message, Chat
|
||||
from pyrogram.enums import ChatType
|
||||
m = Message(id=7, date=datetime.now(timezone.utc), text="hello",
|
||||
chat=Chat(id=-1001, type=ChatType.CHANNEL, username="testchan"))
|
||||
roundtripped = pickle.loads(pickle.dumps(m)) # mimics the pickle cache
|
||||
clone = copy.deepcopy(roundtripped) # what _create_time_based_media_groups does
|
||||
assert clone.id == 7
|
||||
assert clone.chat.username == "testchan"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4.2 — no asyncio in the render path; bulk upsert after render
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_render_path_has_no_asyncio_side_effects():
|
||||
banned = {"create_task", "get_running_loop", "to_thread", "ensure_future"}
|
||||
funcs = [
|
||||
_render_pipeline, _create_time_based_media_groups, _create_messages_groups,
|
||||
_trim_messages_groups, _render_messages_groups,
|
||||
PostParser.process_message, PostParser._generate_html_body,
|
||||
PostParser._generate_html_media, PostParser.generate_html_footer,
|
||||
PostParser._reactions_views_links, PostParser._save_media_file_ids,
|
||||
PostParser._sanitize_html,
|
||||
]
|
||||
for fn in funcs:
|
||||
offenders = _co_names(fn) & banned
|
||||
assert not offenders, f"{fn.__qualname__} references forbidden asyncio names: {offenders}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_media_ids_persisted_via_bulk_upsert(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_bulk(db_path, entries):
|
||||
calls.append(list(entries))
|
||||
|
||||
monkeypatch.setattr(pp_module, "upsert_media_file_ids_bulk_sync", fake_bulk)
|
||||
|
||||
async def fake_get_chat(client, channel):
|
||||
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
|
||||
|
||||
async def fake_get_history(client, channel, limit=20):
|
||||
return [
|
||||
make_message(1, text="just text"),
|
||||
make_message(2, media=MessageMediaType.PHOTO, photo_uid="uid_abc"),
|
||||
]
|
||||
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
|
||||
|
||||
await generate_channel_rss("testchan", client=SimpleNamespace(), limit=10)
|
||||
|
||||
assert len(calls) == 1, "bulk upsert must be called exactly once after render"
|
||||
entries = calls[0]
|
||||
assert len(entries) == 1, "only the photo message contributes a media id"
|
||||
channel, post_id, fid, _ts = entries[0]
|
||||
assert (channel, post_id, fid) == ("testchan", 2, "uid_abc")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_media_file_ids_only_appends(monkeypatch):
|
||||
# Even with a running loop, _save_media_file_ids must not create tasks — just append.
|
||||
parser = PostParser(SimpleNamespace())
|
||||
monkeypatch.setattr(pp_module, "upsert_media_file_ids_bulk_sync",
|
||||
lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not be called directly")))
|
||||
msg = make_message(3, media=MessageMediaType.PHOTO, photo_uid="uid_x")
|
||||
parser._save_media_file_ids(msg)
|
||||
assert parser._pending_media_ids == [("testchan", 3, "uid_x", parser._pending_media_ids[0][3])]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4.1 — raw_message laziness
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_raw_message_lazy_for_feed():
|
||||
parser = PostParser(SimpleNamespace())
|
||||
result = parser.process_message(make_message(10, text="hi"), include_raw=False, sanitize=False)
|
||||
assert "raw_message" not in result
|
||||
|
||||
|
||||
def test_raw_message_present_for_json_and_debug():
|
||||
parser = PostParser(SimpleNamespace())
|
||||
result = parser.process_message(make_message(11, text="hi"), include_raw=True)
|
||||
assert "raw_message" in result
|
||||
assert isinstance(result["raw_message"], str)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_100_message_feed_generates(monkeypatch):
|
||||
async def fake_get_chat(client, channel):
|
||||
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
|
||||
|
||||
async def fake_get_history(client, channel, limit=20):
|
||||
return [make_message(i, text=f"post number {i}") for i in range(1, 101)]
|
||||
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
|
||||
|
||||
rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=100)
|
||||
assert rss.count("<item>") == 100
|
||||
assert "post number 50" in rss
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4.4 — XSS: payload stripped in all four outputs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _assert_clean(html_str, where):
|
||||
assert "<script>" not in html_str, f"{where}: <script> survived"
|
||||
assert "onerror" not in html_str, f"{where}: onerror= survived"
|
||||
assert "javascript:" not in html_str, f"{where}: javascript: survived"
|
||||
|
||||
|
||||
def _make_client_returning(msg):
|
||||
client = SimpleNamespace()
|
||||
|
||||
async def get_messages(channel, post_id):
|
||||
return msg
|
||||
|
||||
client.get_messages = get_messages
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_xss_stripped_in_json():
|
||||
msg = make_message(20, text=XSS_PAYLOAD)
|
||||
parser = PostParser(_make_client_returning(msg))
|
||||
data = await parser.get_post("testchan", 20, "json")
|
||||
_assert_clean(data["html"]["body"], "json body")
|
||||
_assert_clean(data["html"]["footer"], "json footer")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_xss_stripped_in_single_post_html():
|
||||
msg = make_message(21, text=XSS_PAYLOAD)
|
||||
parser = PostParser(_make_client_returning(msg))
|
||||
html_out = await parser.get_post("testchan", 21, "html")
|
||||
_assert_clean(html_out, "single-post html")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_xss_stripped_in_single_post_html_debug():
|
||||
# debug HTML embeds raw_message (str(message)) into a <pre>; it must be html-escaped so
|
||||
# no live tag from the payload survives. The escaped dump legitimately still contains the
|
||||
# inert words "onerror"/"javascript:" as text — what matters is that they are NOT live.
|
||||
msg = make_message(22, text=XSS_PAYLOAD)
|
||||
parser = PostParser(_make_client_returning(msg))
|
||||
html_out = await parser.get_post("testchan", 22, "html", debug=True)
|
||||
|
||||
# 1) The rendered display area (everything before the raw <pre>) is fully sanitized.
|
||||
display = html_out.split('<pre', 1)[0]
|
||||
_assert_clean(display, "single-post debug display")
|
||||
|
||||
# 2) The raw <pre> dump is html-escaped: no live <script> tag anywhere, and the payload
|
||||
# appears only in escaped form (proving html.escape ran).
|
||||
assert "<script>" not in html_out, "debug raw dump left a live <script> tag"
|
||||
assert "<script>" in html_out, "debug raw dump was not html-escaped"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_xss_stripped_in_rss_feed(monkeypatch):
|
||||
async def fake_get_chat(client, channel):
|
||||
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
|
||||
|
||||
async def fake_get_history(client, channel, limit=20):
|
||||
return [make_message(23, text=XSS_PAYLOAD)]
|
||||
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
|
||||
|
||||
rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
|
||||
# The sanitized HTML lives in <content:encoded><![CDATA[...]]></content:encoded>.
|
||||
cdata = re.findall(r"<!\[CDATA\[(.*?)\]\]>", rss, re.DOTALL)
|
||||
assert cdata, "expected CDATA content in RSS"
|
||||
for chunk in cdata:
|
||||
_assert_clean(chunk, "rss content")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_xss_stripped_in_html_feed(monkeypatch):
|
||||
async def fake_get_chat(client, channel):
|
||||
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
|
||||
|
||||
async def fake_get_history(client, channel, limit=20):
|
||||
return [make_message(24, text=XSS_PAYLOAD)]
|
||||
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
|
||||
|
||||
html_feed = await generate_channel_html("testchan", client=SimpleNamespace(), limit=5)
|
||||
_assert_clean(html_feed, "html feed")
|
||||
|
||||
|
||||
def _media_msg_with_payload_caption(mid):
|
||||
# A photo message whose CAPTION carries the payload — this exercises the media
|
||||
# fragment path (_generate_html_media / caption rendering) whose internal per-fragment
|
||||
# sanitize pass was removed in 4.4. The covering pass must still strip it.
|
||||
m = make_message(mid, media=MessageMediaType.PHOTO, photo_uid="pic123")
|
||||
m.caption = _Str(XSS_PAYLOAD)
|
||||
return m
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_xss_in_media_caption_stripped_direct_paths():
|
||||
# json + single-post html go through process_message(sanitize=True) directly.
|
||||
parser = PostParser(_make_client_returning(_media_msg_with_payload_caption(30)))
|
||||
data = await parser.get_post("testchan", 30, "json")
|
||||
_assert_clean(data["html"]["body"], "json body (media caption)")
|
||||
_assert_clean(data["html"]["footer"], "json footer (media caption)")
|
||||
|
||||
parser2 = PostParser(_make_client_returning(_media_msg_with_payload_caption(31)))
|
||||
html_out = await parser2.get_post("testchan", 31, "html")
|
||||
_assert_clean(html_out, "single-post html (media caption)")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_xss_in_media_caption_stripped_in_feeds(monkeypatch):
|
||||
async def fake_get_chat(client, channel):
|
||||
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
|
||||
|
||||
async def fake_get_history(client, channel, limit=20):
|
||||
return [_media_msg_with_payload_caption(32)]
|
||||
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
|
||||
|
||||
rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
|
||||
for chunk in re.findall(r"<!\[CDATA\[(.*?)\]\]>", rss, re.DOTALL):
|
||||
_assert_clean(chunk, "rss content (media caption)")
|
||||
|
||||
async def fake_get_history2(client, channel, limit=20):
|
||||
return [_media_msg_with_payload_caption(33)]
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history2, raising=False)
|
||||
html_feed = await generate_channel_html("testchan", client=SimpleNamespace(), limit=5)
|
||||
_assert_clean(html_feed, "html feed (media caption)")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Review round-1: the new bulk-upsert SQL executed for real (not mocked).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_bulk_upsert_media_file_ids_real_sql(tmp_path):
|
||||
import sqlite3
|
||||
from file_io import upsert_media_file_ids_bulk_sync, init_db_sync
|
||||
|
||||
db = str(tmp_path / "t.db")
|
||||
init_db_sync(db)
|
||||
|
||||
# Empty list is a no-op (no crash).
|
||||
upsert_media_file_ids_bulk_sync(db, [])
|
||||
|
||||
# Multi-row insert.
|
||||
upsert_media_file_ids_bulk_sync(db, [
|
||||
("chA", 1, "fidA", 100.0),
|
||||
("chB", 2, "fidB", 200.0),
|
||||
])
|
||||
conn = sqlite3.connect(db)
|
||||
rows = dict(((c, p, f), a) for c, p, f, a in
|
||||
conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids"))
|
||||
assert rows[("chA", 1, "fidA")] == 100.0
|
||||
assert rows[("chB", 2, "fidB")] == 200.0
|
||||
|
||||
# Re-upsert the SAME key updates `added` (ON CONFLICT ... DO UPDATE SET added=excluded.added).
|
||||
upsert_media_file_ids_bulk_sync(db, [("chA", 1, "fidA", 999.0)])
|
||||
a = conn.execute(
|
||||
"SELECT added FROM media_file_ids WHERE channel='chA' AND post_id=1 AND file_unique_id='fidA'"
|
||||
).fetchone()[0]
|
||||
assert a == 999.0
|
||||
conn.close()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Review round-1: media ids collected before a render exception are still
|
||||
# flushed (the flush is in a finally). Removing the finally must break this.
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_media_flushed_on_render_exception(monkeypatch):
|
||||
async def fake_get_chat(client, channel):
|
||||
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
|
||||
|
||||
async def fake_get_history(client, channel, limit=20):
|
||||
return [make_message(40, text="hi")]
|
||||
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
|
||||
|
||||
# A render pipeline that collects a pending id then raises mid-render.
|
||||
def boom_pipeline(messages, post_parser, *a, **k):
|
||||
post_parser._pending_media_ids.append(("chZ", 9, "fidZ", 1.0))
|
||||
raise RuntimeError("render blew up")
|
||||
|
||||
monkeypatch.setattr(rss_module, "_render_pipeline", boom_pipeline)
|
||||
|
||||
flushed = {}
|
||||
async def fake_bulk(db, entries):
|
||||
flushed["entries"] = list(entries)
|
||||
monkeypatch.setattr("post_parser.upsert_media_file_ids_bulk_sync",
|
||||
lambda db, entries: flushed.__setitem__("entries", list(entries)),
|
||||
raising=False)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
|
||||
|
||||
# The collected id was persisted despite the render exception (flush in finally).
|
||||
assert flushed.get("entries") == [("chZ", 9, "fidZ", 1.0)]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Review round-1 [security]: if the ONLY sanitize pass throws, the feed must
|
||||
# FAIL CLOSED (html.escape the raw content), never emit the raw XSS payload.
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_rss_fails_closed_when_sanitizer_raises(monkeypatch):
|
||||
async def fake_get_chat(client, channel):
|
||||
return SimpleNamespace(title="Test", username="testchan", id=-1001234567890)
|
||||
|
||||
async def fake_get_history(client, channel, limit=20):
|
||||
return [make_message(41, text=XSS_PAYLOAD)]
|
||||
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat", fake_get_chat, raising=False)
|
||||
monkeypatch.setattr("tg_cache.cached_get_chat_history", fake_get_history, raising=False)
|
||||
|
||||
# Force bleach to blow up (e.g. the RecursionError class already seen in prod).
|
||||
# rss_generator imports it as `from bleach import clean as HTMLSanitizer`.
|
||||
def boom(*a, **k):
|
||||
raise RecursionError("bleach exploded")
|
||||
monkeypatch.setattr(rss_module, "HTMLSanitizer", boom, raising=True)
|
||||
|
||||
rss = await generate_channel_rss("testchan", client=SimpleNamespace(), limit=5)
|
||||
chunks = re.findall(r"<!\[CDATA\[(.*?)\]\]>", rss, re.DOTALL)
|
||||
assert chunks, "expected CDATA content"
|
||||
for chunk in chunks:
|
||||
# Fail-closed: the raw payload was html.escaped, so NO live tag survived — every
|
||||
# `<` became `<`. (The letters "javascript:" still appear, but as inert text
|
||||
# inside an escaped "…", not a live href.)
|
||||
assert "<script" not in chunk, "RSS fail-open: raw <script> reached the feed"
|
||||
assert "<img" not in chunk, "RSS fail-open: raw <img onerror> reached the feed"
|
||||
assert "<a " not in chunk, "RSS fail-open: raw <a href> reached the feed"
|
||||
# ...and the escaping actually ran (payload present as escaped text, not dropped).
|
||||
assert "<script>" in chunk, "expected the payload html-escaped, not dropped"
|
||||
@@ -0,0 +1,263 @@
|
||||
# flake8: noqa
|
||||
# pylint: disable=protected-access, missing-function-docstring, missing-class-docstring
|
||||
# pylint: disable=redefined-outer-name, logging-fstring-interpolation, line-too-long
|
||||
# pylance: disable=reportMissingImports, reportMissingModuleSource
|
||||
"""
|
||||
Stage 5 (SQLite access-time batching) tests.
|
||||
|
||||
Covers:
|
||||
- 5.1 accumulator: a /media cache hit records into api_server._access_updates and does NOT
|
||||
call update_media_file_access_sync (no synchronous SQLite write on the hot path).
|
||||
- 5.1 flush: seeding _access_updates + running the flush once writes the accumulated
|
||||
timestamps to a real temp DB (hit -> flush -> value updated in DB), and clears the dict.
|
||||
- 5.1 bulk fn update_media_file_access_bulk_sync: empty no-op, multi-row, and updating an
|
||||
EXISTING row's `added` (WHERE matches on the str channel key).
|
||||
- 5.1 snapshot-then-clear: an update arriving DURING the flush lands in the fresh dict and
|
||||
is not lost.
|
||||
- gotcha: str(channel) key discipline — an int-ish channel on the hot path keys the
|
||||
accumulator (and thus the UPDATE) by the string form.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import sqlite3
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings'])
|
||||
|
||||
import api_server
|
||||
from file_io import init_db_sync, update_media_file_access_bulk_sync
|
||||
|
||||
|
||||
def _fake_plain_message():
|
||||
"""A non-poll, non-video message so download_media_file takes the normal cache flow."""
|
||||
return SimpleNamespace(media=None, video=None, empty=False)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clean_accumulator():
|
||||
"""Each test starts with an empty accumulator and restores it afterwards."""
|
||||
api_server._access_updates = {}
|
||||
yield
|
||||
api_server._access_updates = {}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 5.1 hot path: cache hit records into the accumulator, no synchronous SQLite.
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_hit_records_accumulator_no_sqlite(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
channel, post_id, fid = "testchan", 5, "fidHIT"
|
||||
cache_dir = tmp_path / "data" / "cache" / channel / str(post_id)
|
||||
cache_dir.mkdir(parents=True)
|
||||
(cache_dir / fid).write_bytes(b"cached-bytes")
|
||||
|
||||
async def fake_get(_cid, _pid):
|
||||
return _fake_plain_message()
|
||||
monkeypatch.setattr(api_server.client, "safe_get_messages", fake_get)
|
||||
|
||||
# Spy: the single-row synchronous updater must NOT be called on the hot path.
|
||||
called = []
|
||||
monkeypatch.setattr(api_server, "update_media_file_access_sync",
|
||||
lambda *a, **k: called.append(a))
|
||||
|
||||
path, delete_after = await api_server.download_media_file(channel, post_id, fid)
|
||||
|
||||
assert path == str(cache_dir / fid)
|
||||
assert delete_after is False
|
||||
assert called == [] # no synchronous SQLite access-write happened
|
||||
assert (channel, post_id, fid) in api_server._access_updates
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 5.1 hot path (DoD guard): get_media's pre-semaphore cache-hit — the hottest
|
||||
# changed site, the one this PR exists for — records into the accumulator and
|
||||
# does NO synchronous SQLite access-write. Mirrors the download_media_file spy
|
||||
# so a regression re-introducing a per-hit write into THIS branch goes red.
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_media_pre_semaphore_cache_hit_no_sqlite(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
channel, post_id, fid = "gmchan", 11, "fidGM"
|
||||
cache_dir = tmp_path / "data" / "cache" / channel / str(post_id)
|
||||
cache_dir.mkdir(parents=True)
|
||||
(cache_dir / fid).write_bytes(b"cached-bytes")
|
||||
|
||||
# Bypass the HMAC digest gate and the FileResponse machinery — the test targets
|
||||
# only the access-time write on the pre-semaphore cache-hit branch.
|
||||
monkeypatch.setattr(api_server, "verify_media_digest", lambda *a, **k: True)
|
||||
sentinel = object()
|
||||
async def fake_prepare(cache_path, request=None, media_key=None):
|
||||
return sentinel
|
||||
monkeypatch.setattr(api_server, "prepare_file_response", fake_prepare)
|
||||
|
||||
# Spy: the single-row synchronous updater must NOT be called on the hot path.
|
||||
called = []
|
||||
monkeypatch.setattr(api_server, "update_media_file_access_sync",
|
||||
lambda *a, **k: called.append(a))
|
||||
|
||||
resp = await api_server.get_media(channel, post_id, fid, request=object(), digest="x")
|
||||
|
||||
assert resp is sentinel
|
||||
assert called == [] # no synchronous SQLite access-write on the hottest path
|
||||
assert (channel, post_id, fid) in api_server._access_updates
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# gotcha: str(channel) key discipline on the hot path (int-ish channel).
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_hit_keys_channel_as_str(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
channel_int, post_id, fid = 12345, 7, "fidINT"
|
||||
cache_dir = tmp_path / "data" / "cache" / str(channel_int) / str(post_id)
|
||||
cache_dir.mkdir(parents=True)
|
||||
(cache_dir / fid).write_bytes(b"x")
|
||||
|
||||
async def fake_get(_cid, _pid):
|
||||
return _fake_plain_message()
|
||||
monkeypatch.setattr(api_server.client, "safe_get_messages", fake_get)
|
||||
|
||||
await api_server.download_media_file(channel_int, post_id, fid)
|
||||
|
||||
assert ("12345", post_id, fid) in api_server._access_updates # string form
|
||||
assert (channel_int, post_id, fid) not in api_server._access_updates # never the raw int
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 5.1 flush: hit -> flush -> the accumulated timestamp is written to the DB.
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_writes_accumulated_timestamps(tmp_path, monkeypatch):
|
||||
db = str(tmp_path / "flush.db")
|
||||
init_db_sync(db)
|
||||
# Seed two existing rows with an old timestamp.
|
||||
conn = sqlite3.connect(db)
|
||||
conn.executemany(
|
||||
"INSERT INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?,?,?,?)",
|
||||
[("chA", 1, "fA", 1.0), ("chB", 2, "fB", 2.0)],
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
monkeypatch.setattr(api_server, "DB_PATH", db)
|
||||
api_server._access_updates = {
|
||||
("chA", 1, "fA"): 111.0,
|
||||
("chB", 2, "fB"): 222.0,
|
||||
}
|
||||
|
||||
await api_server._flush_access_updates()
|
||||
|
||||
# Dict cleared after flush.
|
||||
assert api_server._access_updates == {}
|
||||
|
||||
conn = sqlite3.connect(db)
|
||||
rows = dict(((c, p, f), a) for c, p, f, a in
|
||||
conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids"))
|
||||
conn.close()
|
||||
assert rows[("chA", 1, "fA")] == 111.0
|
||||
assert rows[("chB", 2, "fB")] == 222.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_empty_is_noop(monkeypatch):
|
||||
api_server._access_updates = {}
|
||||
# Must not raise and must not touch the threadpool/DB.
|
||||
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync",
|
||||
lambda *a, **k: pytest.fail("bulk sync should not run for an empty batch"))
|
||||
await api_server._flush_access_updates()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 5.1 snapshot-then-clear: an update arriving DURING the flush is not lost.
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_snapshot_then_clear_does_not_lose_late_update(tmp_path, monkeypatch):
|
||||
db = str(tmp_path / "race.db")
|
||||
init_db_sync(db)
|
||||
monkeypatch.setattr(api_server, "DB_PATH", db)
|
||||
|
||||
late_key = ("chLate", 9, "fLate")
|
||||
|
||||
def fake_bulk(_db, entries):
|
||||
# Simulates a cache hit landing WHILE the flush's to_thread runs: because the flush
|
||||
# already replaced the module dict with a fresh one, this write goes into the NEW dict.
|
||||
api_server._access_updates[late_key] = 999.0
|
||||
|
||||
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync", fake_bulk)
|
||||
|
||||
api_server._access_updates = {("chA", 1, "fA"): 111.0}
|
||||
await api_server._flush_access_updates()
|
||||
|
||||
# The snapshot (chA) was flushed and cleared; the late update survives in the new dict.
|
||||
assert api_server._access_updates == {late_key: 999.0}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 5.1 re-queue on failure: a failed bulk write is not lost; setdefault keeps a
|
||||
# fresher write that arrived during the flush.
|
||||
# --------------------------------------------------------------------------- #
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_flush_requeues_batch_without_clobbering_fresh(tmp_path, monkeypatch):
|
||||
db = str(tmp_path / "fail.db")
|
||||
init_db_sync(db)
|
||||
monkeypatch.setattr(api_server, "DB_PATH", db)
|
||||
|
||||
stale_key = ("chStale", 1, "fS") # only in the failing snapshot
|
||||
fresh_key = ("chFresh", 2, "fF") # re-written FRESHER during the flush
|
||||
|
||||
def fake_bulk(_db, _entries):
|
||||
# A newer cache hit for fresh_key lands in the fresh dict WHILE the bulk write runs,
|
||||
# then the write fails. The re-queue must restore stale_key but must NOT overwrite
|
||||
# the newer fresh_key value (setdefault, not assignment).
|
||||
api_server._access_updates[fresh_key] = 999.0
|
||||
raise sqlite3.OperationalError("database is locked")
|
||||
|
||||
monkeypatch.setattr(api_server, "update_media_file_access_bulk_sync", fake_bulk)
|
||||
|
||||
api_server._access_updates = {stale_key: 1.0, fresh_key: 2.0}
|
||||
with pytest.raises(sqlite3.OperationalError):
|
||||
await api_server._flush_access_updates()
|
||||
|
||||
# stale_key restored with its snapshot value; fresh_key keeps the NEWER value (not 2.0).
|
||||
assert api_server._access_updates == {stale_key: 1.0, fresh_key: 999.0}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 5.1 bulk fn: empty no-op, multi-row, and existing-row update via str channel key.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_bulk_access_update_real_sql(tmp_path):
|
||||
db = str(tmp_path / "bulk.db")
|
||||
init_db_sync(db)
|
||||
|
||||
# Empty batch is a no-op (no crash).
|
||||
update_media_file_access_bulk_sync(db, [])
|
||||
|
||||
# Seed existing rows.
|
||||
conn = sqlite3.connect(db)
|
||||
conn.executemany(
|
||||
"INSERT INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?,?,?,?)",
|
||||
[("chA", 1, "fA", 1.0), ("chB", 2, "fB", 2.0)],
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
# Multi-row update of EXISTING rows: WHERE matches on the str channel key.
|
||||
update_media_file_access_bulk_sync(db, [
|
||||
("chA", 1, "fA", 500.0),
|
||||
("chB", 2, "fB", 600.0),
|
||||
])
|
||||
|
||||
conn = sqlite3.connect(db)
|
||||
rows = dict(((c, p, f), a) for c, p, f, a in
|
||||
conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids"))
|
||||
# A row keyed by an int channel does NOT match the string "chA" WHERE (documents the gotcha).
|
||||
n = conn.execute("SELECT COUNT(*) FROM media_file_ids WHERE channel = 1").fetchone()[0]
|
||||
conn.close()
|
||||
assert rows[("chA", 1, "fA")] == 500.0
|
||||
assert rows[("chB", 2, "fB")] == 600.0
|
||||
assert n == 0
|
||||
@@ -0,0 +1,188 @@
|
||||
# flake8: noqa
|
||||
# pylint: disable=protected-access, missing-function-docstring, missing-class-docstring
|
||||
# pylint: disable=redefined-outer-name, logging-fstring-interpolation, line-too-long
|
||||
# pylance: disable=reportMissingImports, reportMissingModuleSource
|
||||
"""
|
||||
Stage 6 (lightweight /ping healthcheck) regression tests.
|
||||
|
||||
Covers:
|
||||
- /ping returns 200 "ok" when connected and the last probe is recent (age < threshold).
|
||||
- /ping returns 503 "degraded" when connected but the last probe is stale (age > threshold).
|
||||
- /ping returns 503 "degraded" when disconnected, regardless of probe age.
|
||||
- /ping returns 200 "ok" on a fresh boot (age is None) while connected — a freshly-started
|
||||
container must NOT be killed before the watchdog's first probe.
|
||||
- ANTI-REGRESSION (the critical invariant): /ping issues ZERO Telegram RPC. A spy on the
|
||||
fake client's get_me / safe_get_messages proves neither is ever called.
|
||||
- TelegramClient.watchdog_last_ok_age(): None when never probed; a positive float afterwards.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
# Add project root to sys.path and mock the config module (same pattern as the other tests).
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings'])
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import api_server
|
||||
from telegram_client import TelegramClient
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Fakes
|
||||
# --------------------------------------------------------------------------- #
|
||||
class _FakeKurigramClient:
|
||||
"""Stands in for TelegramClient.client — exposes is_connected and RPC spies."""
|
||||
def __init__(self, is_connected=True):
|
||||
self.is_connected = is_connected
|
||||
self.get_me_calls = 0
|
||||
|
||||
async def get_me(self):
|
||||
# If /ping ever touches this, the whole point of the endpoint is defeated.
|
||||
self.get_me_calls += 1
|
||||
raise AssertionError("/ping must never call get_me()")
|
||||
|
||||
|
||||
class _FakeTelegramClient:
|
||||
"""Stands in for api_server.client with a controllable probe age + RPC spies."""
|
||||
def __init__(self, age, is_connected=True):
|
||||
self._age = age
|
||||
self.client = _FakeKurigramClient(is_connected=is_connected)
|
||||
self.safe_get_messages_calls = 0
|
||||
|
||||
def watchdog_last_ok_age(self):
|
||||
return self._age
|
||||
|
||||
async def safe_get_messages(self, *args, **kwargs):
|
||||
self.safe_get_messages_calls += 1
|
||||
raise AssertionError("/ping must never call safe_get_messages()")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patch_client(monkeypatch):
|
||||
"""Return a factory that installs a fake api_server.client and yields a TestClient."""
|
||||
def _install(age, is_connected=True):
|
||||
fake = _FakeTelegramClient(age=age, is_connected=is_connected)
|
||||
monkeypatch.setattr(api_server, "client", fake)
|
||||
return fake, TestClient(api_server.app)
|
||||
return _install
|
||||
|
||||
|
||||
THRESHOLD = api_server.Config["tg_ping_unhealthy_after"] # 250 in mock_config
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# /ping endpoint behavior
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_ping_healthy_connected_recent(patch_client):
|
||||
fake, tc = patch_client(age=THRESHOLD - 10, is_connected=True)
|
||||
r = tc.get("/ping")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["status"] == "ok"
|
||||
assert body["connected"] is True
|
||||
assert body["last_probe_age_s"] == round(THRESHOLD - 10, 1)
|
||||
assert body["threshold_s"] == THRESHOLD
|
||||
assert fake.client.get_me_calls == 0
|
||||
|
||||
|
||||
def test_ping_degraded_stale_probe(patch_client):
|
||||
fake, tc = patch_client(age=THRESHOLD + 100, is_connected=True)
|
||||
r = tc.get("/ping")
|
||||
assert r.status_code == 503
|
||||
body = r.json()
|
||||
assert body["status"] == "degraded"
|
||||
assert body["connected"] is True
|
||||
assert fake.client.get_me_calls == 0
|
||||
|
||||
|
||||
def test_ping_watchdog_disabled_stale_age_still_healthy(patch_client, monkeypatch):
|
||||
# With the watchdog OFF, nothing refreshes age (a disconnect-flap restart can stamp it
|
||||
# once, then it only grows). A stale age must NOT drive /ping to 503 on a live connection
|
||||
# — that would spuriously fail the healthcheck and trigger an autoheal restart. So with the
|
||||
# watchdog disabled, /ping is a pure connectivity check: connected + stale age => healthy.
|
||||
monkeypatch.setitem(api_server.Config, "tg_watchdog_enabled", False)
|
||||
fake, tc = patch_client(age=THRESHOLD + 100, is_connected=True)
|
||||
r = tc.get("/ping")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["status"] == "ok"
|
||||
assert body["connected"] is True
|
||||
assert fake.client.get_me_calls == 0
|
||||
|
||||
|
||||
def test_ping_degraded_disconnected(patch_client):
|
||||
# Even with a fresh probe age, a disconnected client is unhealthy.
|
||||
fake, tc = patch_client(age=1.0, is_connected=False)
|
||||
r = tc.get("/ping")
|
||||
assert r.status_code == 503
|
||||
body = r.json()
|
||||
assert body["status"] == "degraded"
|
||||
assert body["connected"] is False
|
||||
|
||||
|
||||
def test_ping_degraded_disconnected_even_when_age_none(patch_client):
|
||||
fake, tc = patch_client(age=None, is_connected=False)
|
||||
r = tc.get("/ping")
|
||||
assert r.status_code == 503
|
||||
assert r.json()["status"] == "degraded"
|
||||
|
||||
|
||||
def test_ping_fresh_boot_age_none_is_healthy(patch_client):
|
||||
# Right after boot the watchdog hasn't probed yet (age None); connected => healthy.
|
||||
fake, tc = patch_client(age=None, is_connected=True)
|
||||
r = tc.get("/ping")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["status"] == "ok"
|
||||
assert body["last_probe_age_s"] is None
|
||||
assert body["threshold_s"] == THRESHOLD
|
||||
|
||||
|
||||
def test_ping_pre_start_connected_none_is_degraded_bool(patch_client):
|
||||
# Before client.start(), Kurigram's is_connected is None. /ping must not 500: it coerces
|
||||
# to a bool, so "connected" is false (never null) and the endpoint reports 503 degraded.
|
||||
fake, tc = patch_client(age=None, is_connected=None)
|
||||
r = tc.get("/ping")
|
||||
assert r.status_code == 503
|
||||
body = r.json()
|
||||
assert body["status"] == "degraded"
|
||||
assert body["connected"] is False # bool, not null
|
||||
assert fake.client.get_me_calls == 0
|
||||
|
||||
|
||||
def test_ping_issues_no_tg_rpc(patch_client):
|
||||
"""The critical invariant: /ping never issues any Telegram RPC in any branch."""
|
||||
for age, connected in [(1.0, True), (THRESHOLD + 500, True), (None, True), (1.0, False)]:
|
||||
fake, tc = patch_client(age=age, is_connected=connected)
|
||||
tc.get("/ping")
|
||||
assert fake.client.get_me_calls == 0, f"get_me called (age={age}, connected={connected})"
|
||||
assert fake.safe_get_messages_calls == 0, f"safe_get_messages called (age={age}, connected={connected})"
|
||||
|
||||
|
||||
def test_ping_route_needs_no_token(patch_client):
|
||||
# /ping is unauthenticated by design (no token path variant); it just works.
|
||||
fake, tc = patch_client(age=1.0, is_connected=True)
|
||||
assert tc.get("/ping").status_code == 200
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# TelegramClient.watchdog_last_ok_age accessor
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_watchdog_last_ok_age_none_when_never_probed():
|
||||
tgc = TelegramClient()
|
||||
assert tgc._wd_last_ok_monotonic is None
|
||||
assert tgc.watchdog_last_ok_age() is None
|
||||
|
||||
|
||||
def test_watchdog_last_ok_age_positive_after_probe():
|
||||
tgc = TelegramClient()
|
||||
tgc._wd_last_ok_monotonic = time.monotonic() - 5
|
||||
age = tgc.watchdog_last_ok_age()
|
||||
assert age is not None
|
||||
assert age >= 5.0
|
||||
# Sanity: a plausible upper bound so we didn't accidentally read the wrong field.
|
||||
assert age < 60.0
|
||||
@@ -0,0 +1,349 @@
|
||||
# flake8: noqa
|
||||
# pylint: disable=protected-access, missing-function-docstring, missing-class-docstring
|
||||
# pylint: disable=redefined-outer-name, logging-fstring-interpolation, line-too-long
|
||||
# pylance: disable=reportMissingImports, reportMissingModuleSource
|
||||
"""
|
||||
Stage 7 — cross-stage END-TO-END integration tests.
|
||||
|
||||
Unlike the per-stage suites (which pin one seam in isolation), these drive the real
|
||||
public entry points (`get_media`, `ping`, the access-time flush) so a regression that
|
||||
only shows up when the stages interact goes red. Every scenario here maps to one of the
|
||||
plan's "Стадия 7 — сквозные ручные сценарии"; the ones that genuinely need a running
|
||||
server + real downloads + lsof (fd-leak counting) are NOT faked here — they stay in
|
||||
docs/stability-verification.md for the operator's prod observation.
|
||||
|
||||
Automated here:
|
||||
- Range semantics at the /media ROUTE level (get_media -> prepare_file_response ->
|
||||
FileResponse driven through real ASGI): bytes=0-99 -> 206, bytes=-100 -> 206,
|
||||
bytes=999999999- -> 416. (Stage 3 pins these on prepare_file_response directly; this
|
||||
adds the end-to-end assertion that get_media's cache-hit branch reaches FileResponse
|
||||
with a live Range still honored — i.e. stages 2+3 wired together.)
|
||||
- /ping (stage 6) stays fast + correct while a deliberately-slow op is in flight, issuing
|
||||
ZERO Telegram RPC — proving the healthcheck is decoupled from the hot/blocked paths.
|
||||
- In-flight dedup + disconnect cleanup (stages 1/2) through the REAL get_media entry
|
||||
point: concurrent requests share one download and the _inflight registry drains; a
|
||||
cancelled request (client disconnect) leaves neither a stuck key nor a hung sibling.
|
||||
- str(channel) access-time key consistency (stage 5) end-to-end: a /media cache hit for
|
||||
an int-ish channel records the str-keyed timestamp, and the flush UPDATE matches the
|
||||
str-keyed DB row (hit -> flush -> DB), the exact affinity gotcha the plan warns about.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import sqlite3
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
# Add project root to sys.path and mock the config module (same pattern as the other tests).
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
sys.modules['config'] = __import__('tests.mock_config', fromlist=['get_settings'])
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import api_server
|
||||
from pyrogram import errors
|
||||
from file_io import init_db_sync
|
||||
|
||||
|
||||
# 2048 deterministic bytes so Range slices are byte-checkable.
|
||||
BODY = bytes(range(256)) * 8
|
||||
SIZE = len(BODY)
|
||||
|
||||
|
||||
def _media_app():
|
||||
"""A bare app that mounts the REAL get_media (and ping) with NO lifespan, so
|
||||
client.start() never runs — a pure cache hit never touches Telegram, and FileResponse
|
||||
computes Range/206/416 at send time, which only happens when driven through ASGI."""
|
||||
app = FastAPI()
|
||||
app.add_api_route("/media/{channel}/{post_id}/{file_unique_id}/{digest}", api_server.get_media, methods=["GET"])
|
||||
app.add_api_route("/media/{channel}/{post_id}/{file_unique_id}", api_server.get_media, methods=["GET"])
|
||||
return app
|
||||
|
||||
|
||||
def _seed_cache(tmp_path, channel, post_id, fid, body=BODY):
|
||||
cache_dir = tmp_path / "data" / "cache" / str(channel) / str(post_id)
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
(cache_dir / fid).write_bytes(body)
|
||||
return cache_dir / fid
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Range semantics at the /media ROUTE level (stages 2 + 3 wired together).
|
||||
# Plan scenario: `curl -H "Range: bytes=0-99" / "bytes=-100" / "bytes=999999999-"`.
|
||||
# Regression caught: any change that makes get_media's cache-hit branch stop reaching
|
||||
# FileResponse (e.g. re-buffering the body, hand-rolling headers, dropping the media_key
|
||||
# path) or that breaks the digest gate wiring — the Range would stop being honored.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_media_route_range_prefix_0_99(monkeypatch, tmp_path):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
_seed_cache(tmp_path, "chan", 3, "fidR")
|
||||
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
|
||||
# Keep the MIME path DB-free; the FileResponse/Range machinery is what we exercise.
|
||||
monkeypatch.setattr(api_server, "get_mime_type_sync", lambda *a, **k: None)
|
||||
monkeypatch.setattr(api_server, "set_mime_type_sync", lambda *a, **k: None)
|
||||
|
||||
c = TestClient(_media_app())
|
||||
r = c.get("/media/chan/3/fidR/anydigest", headers={"Range": "bytes=0-99"})
|
||||
assert r.status_code == 206
|
||||
assert r.headers["content-range"] == f"bytes 0-99/{SIZE}"
|
||||
assert r.headers["content-length"] == "100"
|
||||
assert r.content == BODY[:100]
|
||||
assert r.headers["accept-ranges"] == "bytes"
|
||||
|
||||
|
||||
def test_media_route_range_suffix_last_100(monkeypatch, tmp_path):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
_seed_cache(tmp_path, "chan", 3, "fidS")
|
||||
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
|
||||
monkeypatch.setattr(api_server, "get_mime_type_sync", lambda *a, **k: None)
|
||||
monkeypatch.setattr(api_server, "set_mime_type_sync", lambda *a, **k: None)
|
||||
|
||||
c = TestClient(_media_app())
|
||||
r = c.get("/media/chan/3/fidS/anydigest", headers={"Range": "bytes=-100"})
|
||||
assert r.status_code == 206
|
||||
assert r.headers["content-range"] == f"bytes {SIZE - 100}-{SIZE - 1}/{SIZE}"
|
||||
assert r.headers["content-length"] == "100"
|
||||
assert r.content == BODY[-100:]
|
||||
|
||||
|
||||
def test_media_route_range_unsatisfiable_416(monkeypatch, tmp_path):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
_seed_cache(tmp_path, "chan", 3, "fidU")
|
||||
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
|
||||
monkeypatch.setattr(api_server, "get_mime_type_sync", lambda *a, **k: None)
|
||||
monkeypatch.setattr(api_server, "set_mime_type_sync", lambda *a, **k: None)
|
||||
|
||||
c = TestClient(_media_app())
|
||||
r = c.get("/media/chan/3/fidU/anydigest", headers={"Range": "bytes=999999999-"})
|
||||
assert r.status_code == 416
|
||||
# Starlette's 416 Content-Range is `*/size` (documented stage-3 RFC-7233 delta).
|
||||
assert r.headers["content-range"] == f"*/{SIZE}"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# /ping (stage 6) stays fast + correct while a slow op is in flight, zero TG RPC.
|
||||
# Plan scenario: "Генерация фида на 100+ сообщений + параллельный /ping -> ping < 100 мс".
|
||||
# We model the concurrent slow op as a coroutine parked on an Event that is NEVER set
|
||||
# during the ping, and assert ping resolves while it is still pending AND touches no RPC.
|
||||
# Regression caught: re-coupling /ping to a TG RPC (get_me / safe_get_messages) or to any
|
||||
# awaitable that a blocked hot path could stall — the ping would no longer return promptly
|
||||
# or the spy counts would go non-zero.
|
||||
# --------------------------------------------------------------------------- #
|
||||
class _FakeKurigram:
|
||||
def __init__(self, is_connected=True):
|
||||
self.is_connected = is_connected
|
||||
self.get_me_calls = 0
|
||||
|
||||
async def get_me(self):
|
||||
self.get_me_calls += 1
|
||||
raise AssertionError("/ping must never call get_me()")
|
||||
|
||||
|
||||
class _FakeTelegramClient:
|
||||
def __init__(self, age, is_connected=True):
|
||||
self._age = age
|
||||
self.client = _FakeKurigram(is_connected=is_connected)
|
||||
self.safe_get_messages_calls = 0
|
||||
|
||||
def watchdog_last_ok_age(self):
|
||||
return self._age
|
||||
|
||||
async def safe_get_messages(self, *a, **k):
|
||||
self.safe_get_messages_calls += 1
|
||||
raise AssertionError("/ping must never call safe_get_messages()")
|
||||
|
||||
|
||||
async def test_ping_prompt_and_rpc_free_while_slow_op_pending(monkeypatch):
|
||||
threshold = api_server.Config["tg_ping_unhealthy_after"]
|
||||
fake = _FakeTelegramClient(age=threshold - 10, is_connected=True)
|
||||
monkeypatch.setattr(api_server, "client", fake)
|
||||
|
||||
# A concurrent slow operation (a stand-in for a hung feed/RPC hot path) parked on an
|
||||
# Event we deliberately never set for the duration of the ping.
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def slow_op():
|
||||
await gate.wait()
|
||||
|
||||
slow = asyncio.create_task(slow_op())
|
||||
await asyncio.sleep(0) # let slow_op start and park on the gate
|
||||
|
||||
# The real proof of decoupling: ping() returns under a tight deadline while the slow op
|
||||
# is parked, AND issues zero TG RPC. wait_for reds if ping ever blocks; the spies (which
|
||||
# raise if touched) red if ping recouples to any RPC. `not slow.done()` is only a sanity
|
||||
# check that ping did not somehow drive the parked op — the gate keeps it pending anyway.
|
||||
resp = await asyncio.wait_for(api_server.ping(), timeout=1.0)
|
||||
|
||||
assert resp.status_code == 200 # correct health while connected + fresh
|
||||
assert not slow.done() # sanity: slow op still parked, ping did not await it
|
||||
assert fake.client.get_me_calls == 0 # decoupled: zero TG RPC
|
||||
assert fake.safe_get_messages_calls == 0
|
||||
|
||||
gate.set()
|
||||
await slow
|
||||
|
||||
|
||||
async def test_ping_reports_degraded_promptly_while_slow_op_pending(monkeypatch):
|
||||
threshold = api_server.Config["tg_ping_unhealthy_after"]
|
||||
fake = _FakeTelegramClient(age=threshold + 100, is_connected=True) # stale probe
|
||||
monkeypatch.setattr(api_server, "client", fake)
|
||||
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def slow_op():
|
||||
await gate.wait()
|
||||
|
||||
slow = asyncio.create_task(slow_op())
|
||||
await asyncio.sleep(0)
|
||||
|
||||
resp = await asyncio.wait_for(api_server.ping(), timeout=1.0)
|
||||
|
||||
assert resp.status_code == 503 # stale watchdog probe -> degraded, still instant
|
||||
assert not slow.done() # sanity: slow op still parked, ping did not await it
|
||||
assert fake.client.get_me_calls == 0
|
||||
assert fake.safe_get_messages_calls == 0
|
||||
|
||||
gate.set()
|
||||
await slow
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# In-flight dedup + disconnect cleanup (stages 1/2) through the REAL get_media entry.
|
||||
# Plan scenario: "Параллельные запросы одного большого видео -> нет частичной отдачи" and
|
||||
# "Отключение клиента на середине стрима -> нет утечки тасков/фд".
|
||||
# The pure-unit stage-2 tests pin _download_deduped directly; these drive get_media so the
|
||||
# HTTP semaphore + dedup registry + serve path are proven wired together.
|
||||
# Regression caught: moving the download back into the request coroutine (so a disconnect
|
||||
# cancels it), or dropping the finally-pop, would leave a stuck _inflight key / hung sibling.
|
||||
# --------------------------------------------------------------------------- #
|
||||
async def test_get_media_concurrent_shares_one_download_and_drains_registry(monkeypatch, tmp_path):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
api_server._inflight.clear()
|
||||
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
|
||||
# Serve step is not under test here; keep it to a sentinel so we assert on dedup + registry.
|
||||
sentinel = object()
|
||||
|
||||
async def fake_prepare(file_path, request=None, delete_after=False, media_key=None):
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(api_server, "prepare_file_response", fake_prepare)
|
||||
|
||||
calls = []
|
||||
|
||||
async def slow_dl(channel, post_id, fid):
|
||||
calls.append(fid)
|
||||
await asyncio.sleep(0.05) # real overlap window for the two requests
|
||||
return (f"/final/{fid}", False)
|
||||
|
||||
monkeypatch.setattr(api_server, "download_media_file", slow_dl)
|
||||
|
||||
req = SimpleNamespace(headers={})
|
||||
t1 = asyncio.create_task(api_server.get_media("chan", 9, "vfid", request=req, digest="x"))
|
||||
t2 = asyncio.create_task(api_server.get_media("chan", 9, "vfid", request=req, digest="x"))
|
||||
r1, r2 = await asyncio.gather(t1, t2)
|
||||
|
||||
assert r1 is sentinel and r2 is sentinel
|
||||
assert calls == ["vfid"] # exactly ONE real download, shared by both requests
|
||||
assert not api_server._inflight # registry drained — no forever-busy key
|
||||
|
||||
|
||||
async def test_get_media_request_cancel_does_not_stick_registry_or_hang_sibling(monkeypatch, tmp_path):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
api_server._inflight.clear()
|
||||
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
|
||||
sentinel = object()
|
||||
|
||||
async def fake_prepare(file_path, request=None, delete_after=False, media_key=None):
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(api_server, "prepare_file_response", fake_prepare)
|
||||
|
||||
gate = asyncio.Event()
|
||||
calls = []
|
||||
|
||||
async def held_dl(channel, post_id, fid):
|
||||
calls.append(fid)
|
||||
await gate.wait() # hold the shared download open until we release it
|
||||
return (f"/final/{fid}", False)
|
||||
|
||||
monkeypatch.setattr(api_server, "download_media_file", held_dl)
|
||||
|
||||
req = SimpleNamespace(headers={})
|
||||
t1 = asyncio.create_task(api_server.get_media("chan", 9, "cfid", request=req, digest="x"))
|
||||
await asyncio.sleep(0.02) # t1 registers the future + starts the detached download
|
||||
t2 = asyncio.create_task(api_server.get_media("chan", 9, "cfid", request=req, digest="x"))
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
# First client disconnects: its request coroutine is cancelled mid-wait.
|
||||
t1.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await t1
|
||||
|
||||
# The detached download is unaffected; releasing it resolves the surviving request.
|
||||
gate.set()
|
||||
r2 = await asyncio.wait_for(t2, timeout=2.0)
|
||||
assert r2 is sentinel
|
||||
assert calls == ["cfid"] # download ran exactly once (not restarted)
|
||||
assert not api_server._inflight # registry drained despite the disconnect
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# str(channel) access-time key consistency (stage 5) END-TO-END: hit -> flush -> DB.
|
||||
# Plan gotcha #9: "Ключи SQLite: channel всегда str(...)"; if the key the hot path RECORDS
|
||||
# and the key the flush UPDATEs by ever diverge, the timestamp silently stops updating and
|
||||
# the file eventually falls out of the cache. Stage 5 pins hit and flush separately; this
|
||||
# stitches them: the /media cache-hit records a (str-channel) key, and the flush's bulk
|
||||
# UPDATE must land on that exact DB row. (get_media's route channel is already a str, so the
|
||||
# str() there is a defensive no-op; the live int/str hazard is in download_media_file, pinned
|
||||
# by stage-5's test_cache_hit_keys_channel_as_str — this test covers the get_media+flush leg.)
|
||||
# Regression caught: any accumulator-key vs UPDATE-WHERE inconsistency (e.g. a transposed
|
||||
# key-column order in the bulk SQL, verified to turn this test red) leaves `added` stale.
|
||||
# --------------------------------------------------------------------------- #
|
||||
async def test_media_cache_hit_flush_updates_str_keyed_db_row(monkeypatch, tmp_path):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
api_server._access_updates = {}
|
||||
|
||||
channel, post_id, fid = "12345", 7, "fidINT" # int-ish channel, passed as the route str
|
||||
_seed_cache(tmp_path, channel, post_id, fid)
|
||||
|
||||
db = str(tmp_path / "access.db")
|
||||
init_db_sync(db)
|
||||
# Seed the row keyed by the STRING channel with an old access time.
|
||||
conn = sqlite3.connect(db)
|
||||
conn.execute(
|
||||
"INSERT INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?,?,?,?)",
|
||||
(channel, post_id, fid, 1.0),
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
monkeypatch.setattr(api_server, "DB_PATH", db)
|
||||
|
||||
monkeypatch.setattr(api_server, "verify_media_digest", lambda url, digest: True)
|
||||
sentinel = object()
|
||||
|
||||
async def fake_prepare(file_path, request=None, media_key=None):
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(api_server, "prepare_file_response", fake_prepare)
|
||||
|
||||
before = time.time()
|
||||
resp = await api_server.get_media(channel, post_id, fid, request=SimpleNamespace(headers={}), digest="x")
|
||||
assert resp is sentinel
|
||||
|
||||
# Hot path recorded the str-keyed access time (never the raw int form).
|
||||
assert (channel, post_id, fid) in api_server._access_updates
|
||||
|
||||
# Flush the accumulator; the bulk UPDATE must match the str-keyed row and refresh `added`.
|
||||
await api_server._flush_access_updates()
|
||||
assert api_server._access_updates == {}
|
||||
|
||||
conn = sqlite3.connect(db)
|
||||
added = conn.execute(
|
||||
"SELECT added FROM media_file_ids WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
|
||||
(channel, post_id, fid),
|
||||
).fetchone()[0]
|
||||
conn.close()
|
||||
assert added >= before # refreshed from the stale 1.0 to ~now via the str-keyed WHERE
|
||||
Reference in New Issue
Block a user