Merge pull request 'fix(stability): стадия 4 — гигиена event loop (рендер фида в поток, ленивый raw_message, один sanitize на выход)' (#12) from fix/stage-4-eventloop into fix/stage-3-fileresponse

This commit was merged in pull request #12.
This commit is contained in:
2026-07-05 16:58:31 +03:00
13 changed files with 1773 additions and 92 deletions
+116 -22
View File
@@ -33,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
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
@@ -139,6 +140,59 @@ 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"])
@@ -161,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:
@@ -172,6 +228,17 @@ 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)
@@ -538,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)
@@ -1003,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:
@@ -1074,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))
+11
View File
@@ -123,6 +123,17 @@ 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).
+8 -3
View File
@@ -53,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
+148
View File
@@ -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
View File
@@ -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
View File
@@ -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('&nbsp;|&nbsp;'.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)}")
+81 -25
View File
@@ -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")
+11
View File
@@ -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):
+1
View File
@@ -33,6 +33,7 @@ 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,
+471
View File
@@ -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 "&lt;script&gt;" 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 `&lt;`. (The letters "javascript:" still appear, but as inert text
# inside an escaped &quot;…&quot;, 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 "&lt;script&gt;" in chunk, "expected the payload html-escaped, not dropped"
+263
View File
@@ -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
+188
View File
@@ -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
+349
View File
@@ -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