pyrogram-bridge/tests/test_media_selfheal.py

# flake8: noqa
# pylint: disable=protected-access, missing-function-docstring, redefined-outer-name, line-too-long
# pylance: disable=reportMissingImports, reportMissingModuleSource
"""
Regression tests for the media self-healing fix (post 'static refactor' outage):

Root cause recap: Kurigram serializes downloads through a single get_file slot; a
zombie media-DC connection makes every download time out while the main-DC watchdog
stays green, so downloads jam forever. The fix adds (a) a negative-cache backoff so a
repeatedly-failing file is not hammered, and (b) a consecutive-timeout streak that
reuses the existing in-process restart to rebuild the media connection.
"""
import asyncio

import pytest

import api_server
from telegram_client import TelegramClient


# --------------------------------------------------------------------------- #
# Negative-cache backoff helpers
# --------------------------------------------------------------------------- #
def test_backoff_arms_and_clears():
    key = ("selfheal_chan", 1, "fid_a")
    api_server._download_failures.pop(key, None)
    assert api_server._download_backoff_remaining(key) == 0.0  # never failed -> allowed
    api_server._record_download_failure(key)
    assert api_server._download_backoff_remaining(key) > 0.0   # armed -> blocked
    api_server._clear_download_failure(key)
    assert api_server._download_backoff_remaining(key) == 0.0  # recovered -> allowed


def test_backoff_failure_counter_increments():
    key = ("selfheal_chan", 2, "fid_b")
    api_server._download_failures.pop(key, None)
    api_server._record_download_failure(key)
    first_fails = api_server._download_failures[key][0]
    api_server._record_download_failure(key)
    second_fails = api_server._download_failures[key][0]
    assert first_fails == 1
    assert second_fails == 2  # consecutive-failure counter grows -> longer backoff
    api_server._clear_download_failure(key)


# --------------------------------------------------------------------------- #
# Download-timeout streak -> single media-connection restart
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_timeout_streak_triggers_single_restart(monkeypatch):
    c = TelegramClient()
    calls = []

    async def fake_restart(reason: str = "unspecified"):
        calls.append(reason)

    monkeypatch.setattr(c, "_restart_client", fake_restart)
    threshold = c.media_timeout_restart_threshold

    # One short of the threshold: no restart scheduled yet.
    for _ in range(threshold - 1):
        c.note_download_timeout()
    assert c._media_recovery_task is None
    assert c._download_timeout_streak == threshold - 1

    # The threshold-th timeout schedules exactly one restart and resets the streak.
    c.note_download_timeout()
    assert c._download_timeout_streak == 0
    assert c._media_recovery_task is not None
    await c._media_recovery_task
    assert calls == ["media download timeout streak"]


@pytest.mark.asyncio
async def test_success_resets_streak():
    c = TelegramClient()
    c.note_download_timeout()
    c.note_download_timeout()
    assert c._download_timeout_streak == 2
    c.note_download_ok()
    assert c._download_timeout_streak == 0


@pytest.mark.asyncio
async def test_no_restart_while_already_restarting(monkeypatch):
    c = TelegramClient()
    calls = []

    async def fake_restart(reason: str = "unspecified"):
        calls.append(reason)

    monkeypatch.setattr(c, "_restart_client", fake_restart)
    c._restarting = True  # a restart is already underway

    for _ in range(c.media_timeout_restart_threshold):
        c.note_download_timeout()

    # Streak reached the threshold but no new restart is scheduled during a restart.
    assert c._media_recovery_task is None
    assert calls == []


# --------------------------------------------------------------------------- #
# Integration: jam -> auto-recovery -> negative cache cleared (Fix 1)
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_hung_download_times_out_and_frees_slot(monkeypatch):
    """A hung download must be cancelled by asyncio.wait_for, freeing its transmission slot
    for the next download (and counting toward the restart streak)."""
    c = TelegramClient()

    # A single-permit semaphore stands in for Pyrogram's get_file transmission slot: the mock
    # holds it for the whole (never-completing) download and releases it in the async-with
    # __aexit__, which runs when wait_for cancels the coroutine.
    slot = asyncio.Semaphore(1)
    entered = asyncio.Event()

    async def hung_download(file_id, file_name=None):
        async with slot:
            entered.set()
            await asyncio.sleep(3600)  # never completes within the test's timeout

    monkeypatch.setattr(c.client, "download_media", hung_download)

    with pytest.raises(asyncio.TimeoutError):
        await c.safe_download_media("fid", "/tmp/x", max_retries=1, timeout=0.05)

    # (a) the download actually started and its slot was released by the cancellation, so a
    # fresh acquire for the "next" download succeeds immediately.
    assert entered.is_set()
    await asyncio.wait_for(slot.acquire(), timeout=1.0)
    slot.release()
    # The timeout was counted toward the restart streak.
    assert c._download_timeout_streak == 1


@pytest.mark.asyncio
async def test_streak_escalates_once_and_verified_restart_clears_cache(monkeypatch):
    """Full self-heal chain: a timeout streak reaching the threshold escalates to EXACTLY ONE
    restart, and the verified restart CLEARS the download negative cache (Fix 1)."""
    c = TelegramClient()

    # Pre-arm the negative cache with a persistently-failing file.
    key = ("selfheal_chan", 7, "fid_jam")
    api_server._download_failures.pop(key, None)
    api_server._record_download_failure(key)
    assert api_server._download_backoff_remaining(key) > 0.0

    # Wire the REAL self-heal hook and drive a REAL _restart_client with only the network
    # layer mocked, so the whole chain (streak -> one restart -> verify ok -> cache clear) runs.
    c.set_restart_callback(api_server._clear_all_download_failures)
    monkeypatch.setattr(c, "_start_watchdog", lambda: None)
    monkeypatch.setattr(c.client, "is_connected", True)

    restart_calls = []

    async def fake_client_restart():
        restart_calls.append(True)

    async def fake_get_me():
        return type("Me", (), {"id": 42})()

    monkeypatch.setattr(c.client, "restart", fake_client_restart)
    monkeypatch.setattr(c.client, "get_me", fake_get_me)

    threshold = c.media_timeout_restart_threshold
    for _ in range(threshold - 1):
        c.note_download_timeout()
    assert c._media_recovery_task is None  # one short of threshold: nothing scheduled yet

    c.note_download_timeout()  # threshold-th timeout schedules exactly one restart
    assert c._media_recovery_task is not None
    await c._media_recovery_task

    # (b) escalated to EXACTLY ONE underlying client.restart().
    assert restart_calls == [True]
    # (c) the verified restart cleared the negative cache -> the file loads again immediately.
    assert api_server._download_backoff_remaining(key) == 0.0
    assert len(api_server._download_failures) == 0


@pytest.mark.asyncio
async def test_failed_verify_does_not_clear_cache(monkeypatch):
    """The cache-clear hook must fire ONLY on a verified restart: if verify_get_me fails, the
    backoff must survive (a still-broken media DC should not drop the protective backoff)."""
    c = TelegramClient()
    key = ("selfheal_chan", 8, "fid_still_broken")
    api_server._download_failures.pop(key, None)
    api_server._record_download_failure(key)

    c.set_restart_callback(api_server._clear_all_download_failures)
    monkeypatch.setattr(c, "_start_watchdog", lambda: None)
    monkeypatch.setattr(c.client, "is_connected", True)

    async def fake_client_restart():
        return None

    async def failing_get_me():
        raise RuntimeError("media DC still dead")

    monkeypatch.setattr(c.client, "restart", fake_client_restart)
    monkeypatch.setattr(c.client, "get_me", failing_get_me)

    await c._restart_client(reason="test failed verify")

    # verify_get_me failed -> callback NOT fired -> backoff still armed.
    assert api_server._download_backoff_remaining(key) > 0.0
    api_server._clear_download_failure(key)


# --------------------------------------------------------------------------- #
# Cross-path negative-cache key consistency
# --------------------------------------------------------------------------- #
def test_download_key_consistent_across_paths():
    """The negative-cache key must be byte-identical across every producer/consumer path so a
    failure recorded on one path is seen by the guard on another. Reproduce each path's key
    expression for the same logical file and assert they coincide, then verify end-to-end that
    a worker-recorded failure is visible to the get_media backoff check."""
    channel_from_db = "durov"   # background_download_worker / download_new_files: str(channel)
    post_id_from_db = 123       # ...                                              int(post_id)
    fid = "AgADfid"

    # get_media receives the raw (possibly differently-cased) URL channel, canonicalizes it,
    # and uses the int path param post_id.
    raw_url_channel = "Durov"
    fs_channel = api_server.canonical_channel_key(raw_url_channel)
    get_media_key = (fs_channel, 123, fid)

    worker_key = (str(channel_from_db), int(post_id_from_db), fid)      # background_download_worker
    new_files_key = (str(channel_from_db), int(post_id_from_db), fid)   # download_new_files
    deduped_key = (str(fs_channel), 123, fid)                          # _download_deduped(fs_channel, post_id, fid)

    assert worker_key == new_files_key == get_media_key == deduped_key

    # End-to-end: a failure recorded by the worker path is seen by the get_media guard.
    api_server._download_failures.pop(worker_key, None)
    api_server._record_download_failure(worker_key)
    try:
        assert api_server._download_backoff_remaining(get_media_key) > 0.0
    finally:
        api_server._clear_download_failure(worker_key)


# --------------------------------------------------------------------------- #
# Fix 2: negative cache is LRU-bounded
# --------------------------------------------------------------------------- #
def test_download_failures_lru_bounded(monkeypatch):
    """The negative cache must not grow unbounded (permanently-404 files leak entries).
    Once over the cap the oldest entry is evicted."""
    monkeypatch.setattr(api_server, "_DOWNLOAD_FAILURES_MAX", 3)
    api_server._download_failures.clear()
    try:
        for i in range(5):
            api_server._record_download_failure(("chan", i, "fid"))
        assert len(api_server._download_failures) == 3
        # Oldest two (i=0,1) evicted; newest three retained.
        remaining = list(api_server._download_failures.keys())
        assert remaining == [("chan", 2, "fid"), ("chan", 3, "fid"), ("chan", 4, "fid")]
    finally:
        api_server._download_failures.clear()


# --------------------------------------------------------------------------- #
# Backoff is capped at _DOWNLOAD_BACKOFF_MAX (min() cap at api_server.py:163)
# --------------------------------------------------------------------------- #
def test_backoff_capped_at_max():
    """The exponential backoff grows as BASE * 2**(fails-1), but must never exceed
    _DOWNLOAD_BACKOFF_MAX. With BASE=60s and MAX=600s the raw exponential overtakes the cap
    at the 5th failure (60*2**4 = 960s > 600s), so after that many consecutive failures the
    effective backoff must stay pinned at the cap. This test FAILS if the min(MAX, ...) cap
    is removed (the raw exponential would then blow past 600s)."""
    key = ("selfheal_chan", 99, "fid_cap")
    api_server._download_failures.pop(key, None)

    # Record well past the point where the raw exponential exceeds the cap: 8 failures ->
    # raw 60*2**7 = 7680s, an order of magnitude above the 600s cap.
    for _ in range(8):
        api_server._record_download_failure(key)

    assert api_server._download_failures[key][0] == 8  # counter really climbed that high
    remaining = api_server._download_backoff_remaining(key)
    # The cap holds: remaining is bounded by MAX (with a tiny slack for monotonic drift since
    # retry_not_before was stamped). Without the min() cap this would be ~7680s.
    assert remaining <= api_server._DOWNLOAD_BACKOFF_MAX
    assert remaining > api_server._DOWNLOAD_BACKOFF_MAX - 5  # and it IS pinned near the cap

    api_server._clear_download_failure(key)


# --------------------------------------------------------------------------- #
# Restart-verified callback is best-effort (try/except at telegram_client.py:250-254)
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_restart_callback_error_does_not_abort(monkeypatch, caplog):
    """A registered restart-verified callback that RAISES must not abort recovery: the error is
    swallowed + logged, and the restart still completes normally (watchdog re-armed, no SIGTERM
    fallback, _restarting reset). This test FAILS if the try/except around the callback is
    removed — the exception would then propagate to the outer handler, skip the watchdog re-arm
    and trigger the process-restart fallback instead."""
    c = TelegramClient()

    def exploding_callback():
        raise RuntimeError("boom in verified-restart callback")

    c.set_restart_callback(exploding_callback)

    # Drive a VERIFIED restart: connected client, restart + get_me both succeed so verify_ok is
    # True and the callback fires.
    monkeypatch.setattr(c.client, "is_connected", True)

    async def fake_client_restart():
        return None

    async def fake_get_me():
        return type("Me", (), {"id": 42})()

    monkeypatch.setattr(c.client, "restart", fake_client_restart)
    monkeypatch.setattr(c.client, "get_me", fake_get_me)

    # Observe the success path (watchdog re-arm) vs the failure path (SIGTERM fallback) without
    # actually killing the test process.
    watchdog_rearmed = []
    monkeypatch.setattr(c, "_start_watchdog", lambda: watchdog_rearmed.append(True))
    sigterm_calls = []
    monkeypatch.setattr(c, "_restart_app", lambda: sigterm_calls.append(True))

    with caplog.at_level("WARNING"):
        await c._restart_client(reason="test callback raises")

    # (a) Recovery completed on the SUCCESS path despite the callback raising: watchdog re-armed,
    # no process-restart fallback, restart flag cleared.
    assert watchdog_rearmed == [True]
    assert sigterm_calls == []
    assert c._restarting is False
    # (b) The callback error was logged (not silently dropped).
    assert any(
        "callback raised" in r.getMessage() and "RuntimeError" in r.getMessage()
        for r in caplog.records
    )