import asyncio import random from datetime import datetime, timezone import pytest from bot import fasttransfer as ft from bot.errors import Cancelled @pytest.fixture(autouse=True) def fast_retries(monkeypatch): monkeypatch.setattr(ft, "RETRY_BASE_DELAY", 0.0) @pytest.fixture(autouse=True) def fresh_stats(monkeypatch): monkeypatch.setattr(ft, "stats", ft.TransferStats()) class FakeTransport: """Serves bytes from memory and records uploaded parts.""" def __init__(self, blob=b"", *, fail_first=0, jitter=True, short_once=False, fatal=None, store=None, hang_after=None): self.blob = blob self.fail_first = fail_first self.jitter = jitter self.short_once = short_once self.fatal = fatal self.calls = 0 self.closed = False self.store = store if store is not None else {} self.hang_after = hang_after self.requests = [] async def _tick(self): self.calls += 1 if self.jitter: await asyncio.sleep(random.random() * 0.003) if self.fatal: raise self.fatal if self.calls <= self.fail_first: raise ConnectionError("boom") if self.hang_after is not None and self.calls > self.hang_after: await asyncio.sleep(3600) async def get_file(self, location, offset, limit): await self._tick() self.requests.append((offset, limit)) assert offset % 4096 == 0 and limit % 4096 == 0 assert (offset % (1024 * 1024)) + limit <= 1024 * 1024 # Telegram: never cross a 1 MiB boundary data = self.blob[offset: offset + limit] if self.short_once and offset > 0 and len(data) == limit: self.short_once = False return data[: limit // 2] return data async def save_part(self, file_id, part, total_parts, data): await self._tick() self.store[part] = (file_id, total_parts, bytes(data)) async def close(self): self.closed = True def blob_of(n, seed=1): return random.Random(seed).randbytes(n) @pytest.mark.parametrize("size", [0, 1, 4095, 1024 * 1024 - 1, 1024 * 1024, 1024 * 1024 + 1, 5 * 1024 * 1024 + 12345]) async def test_download_matches_source(tmp_path, size): blob = blob_of(size) transports = [FakeTransport(blob) for _ in range(4)] seen = [] n = await ft.download_parts(transports, "loc", size, tmp_path / "out.bin", progress=lambda c, t: seen.append((c, t))) assert n == size assert (tmp_path / "out.bin").read_bytes() == blob if size: assert seen[-1] == (size, size) assert all(t == size for _c, t in seen) assert max(c for c, _t in seen) == size async def test_download_work_is_shared_between_connections(tmp_path): size = 8 * 1024 * 1024 blob = blob_of(size) transports = [FakeTransport(blob) for _ in range(4)] await ft.download_parts(transports, "loc", size, tmp_path / "o", progress=None) assert sum(t.calls for t in transports) == 8 assert sum(1 for t in transports if t.calls) > 1 # more than one connection did work offsets = sorted(o for t in transports for o, _ in t.requests) assert offsets == [i * 1024 * 1024 for i in range(8)] # every part exactly once async def test_download_retries_transient_failures(tmp_path): size = 3 * 1024 * 1024 + 7 blob = blob_of(size) transports = [FakeTransport(blob, fail_first=2), FakeTransport(blob, fail_first=3)] await ft.download_parts(transports, "loc", size, tmp_path / "o") assert (tmp_path / "o").read_bytes() == blob async def test_download_retries_short_reads(tmp_path): size = 4 * 1024 * 1024 blob = blob_of(size) transports = [FakeTransport(blob, short_once=True)] await ft.download_parts(transports, "loc", size, tmp_path / "o") assert (tmp_path / "o").read_bytes() == blob async def test_download_gives_up_after_max_retries(tmp_path): blob = blob_of(2 * 1024 * 1024) transport = FakeTransport(blob, fail_first=10_000) with pytest.raises(ConnectionError): await ft.download_parts([transport], "loc", len(blob), tmp_path / "o") assert transport.calls == ft.MAX_RETRIES + 1 async def test_download_fatal_error_is_not_retried(tmp_path): class FileReferenceExpiredError(Exception): pass transport = FakeTransport(b"x" * 10, fatal=FileReferenceExpiredError("expired")) with pytest.raises(FileReferenceExpiredError): await ft.download_parts([transport], "loc", 10, tmp_path / "o") assert transport.calls == 1 async def test_download_cancel_stops_workers(tmp_path): size = 64 * 1024 * 1024 blob = blob_of(1024) cancel = asyncio.Event() class SlowTransport(FakeTransport): async def get_file(self, location, offset, limit): await asyncio.sleep(0.01) if offset >= 3 * 1024 * 1024: cancel.set() return b"\0" * limit transports = [SlowTransport(blob) for _ in range(4)] with pytest.raises(Cancelled): await ft.download_parts(transports, "loc", size, tmp_path / "o", cancel=cancel) async def test_download_hung_connection_times_out_and_retries(tmp_path, monkeypatch): monkeypatch.setattr(ft, "PART_TIMEOUT", 0.05) blob = blob_of(2 * 1024 * 1024) hanging = FakeTransport(blob, hang_after=0, jitter=False) # every request hangs with pytest.raises(asyncio.TimeoutError): await ft.download_parts([hanging], "loc", len(blob), tmp_path / "o") assert hanging.calls == ft.MAX_RETRIES + 1 async def test_download_partial_max_bytes(tmp_path): size = 6 * 1024 * 1024 blob = blob_of(size) limit = int(2.5 * 1024 * 1024) n = await ft.download_parts([FakeTransport(blob), FakeTransport(blob)], "loc", size, tmp_path / "o", max_bytes=limit) assert n == limit assert (tmp_path / "o").read_bytes() == blob[:limit] @pytest.mark.parametrize("size", [1, ft.UPLOAD_PART - 1, ft.UPLOAD_PART, ft.UPLOAD_PART + 1, 3 * 1024 * 1024 + 321]) async def test_upload_parts_reassemble(tmp_path, size): blob = blob_of(size, seed=7) src = tmp_path / "in.bin" src.write_bytes(blob) store = {} transports = [FakeTransport(store=store) for _ in range(3)] seen = [] n = await ft.upload_parts(transports, src, size, 1234, progress=lambda c, t: seen.append((c, t))) expected_parts = -(-size // ft.UPLOAD_PART) assert n == expected_parts assert sorted(store) == list(range(expected_parts)) assert all(fid == 1234 and total == expected_parts for fid, total, _d in store.values()) sizes = [len(store[i][2]) for i in range(expected_parts)] assert all(s == ft.UPLOAD_PART for s in sizes[:-1]) # every part but the last is full size assert b"".join(store[i][2] for i in range(expected_parts)) == blob assert seen[-1] == (size, size) async def test_upload_retries_and_cancel(tmp_path): blob = blob_of(2 * 1024 * 1024) src = tmp_path / "in.bin" src.write_bytes(blob) store = {} await ft.upload_parts([FakeTransport(store=store, fail_first=2), FakeTransport(store=store)], src, len(blob), 1) assert b"".join(store[i][2] for i in sorted(store)) == blob cancel = asyncio.Event() cancel.set() with pytest.raises(Cancelled): await ft.upload_parts([FakeTransport(store={})], src, len(blob), 1, cancel=cancel) async def test_upload_detects_truncated_source(tmp_path): src = tmp_path / "in.bin" src.write_bytes(b"a" * 100) with pytest.raises(IOError): await ft.upload_parts([FakeTransport()], src, 5000, 1) # claims more bytes than the file holds async def test_close_all_swallows_errors(): class Bad: async def close(self): raise RuntimeError("x") good = FakeTransport() await ft.close_all([Bad(), good]) assert good.closed # --------------------------------------------------------------------------- # wrappers / fallback # --------------------------------------------------------------------------- class FakeClient: def __init__(self): self.downloaded = [] self.uploaded = [] async def download_media(self, message, file, progress_callback): progress_callback(5, 10) with open(file, "wb") as fh: fh.write(b"stock") self.downloaded.append(file) async def upload_file(self, path, file_name, progress_callback): progress_callback(1, 1) self.uploaded.append((path, file_name)) return ("stock-input-file", file_name) async def iter_download(self, doc, request_size, limit): for _ in range(limit): yield b"z" * request_size class Msg: def __init__(self, document="doc"): self.document = document async def test_download_falls_back_to_stock_when_fast_path_breaks(tmp_path, monkeypatch): async def broken(*a, **k): raise RuntimeError("private API changed") monkeypatch.setattr(ft, "fast_download", broken) client = FakeClient() dest = tmp_path / "x.bin" dest.write_bytes(b"partial garbage") await ft.download(client, Msg(), dest, fast=True) assert dest.read_bytes() == b"stock" assert (ft.stats.fast_down, ft.stats.standard_down, ft.stats.fallbacks) == (0, 1, 1) assert ft.stats.last_fallback == "download: RuntimeError: private API changed" async def test_download_cancel_is_not_swallowed_by_fallback(tmp_path, monkeypatch): async def cancelled(*a, **k): raise Cancelled() monkeypatch.setattr(ft, "fast_download", cancelled) client = FakeClient() with pytest.raises(Cancelled): await ft.download(client, Msg(), tmp_path / "x", fast=True) assert client.downloaded == [] assert ft.stats == ft.TransferStats() # cancelling is not a failure of the fast path async def test_a_fallback_is_logged_as_an_error_with_its_traceback_so_the_log_channel_hears_of_it(tmp_path, monkeypatch, caplog): async def broken(*a, **k): raise RuntimeError("private API changed") monkeypatch.setattr(ft, "fast_download", broken) monkeypatch.setattr(ft, "fast_upload", broken) src = tmp_path / "f.bin" src.write_bytes(b"x") with caplog.at_level("WARNING"): await ft.download(FakeClient(), Msg(), tmp_path / "d.bin", fast=True) await ft.upload(FakeClient(), src, "n.mkv", fast=True) reports = [r for r in caplog.records if r.levelname == "ERROR" and r.exc_info] assert [r.getMessage().split(" failed")[0] for r in reports] == ["fast download", "fast upload"] assert all("Telethon's standard" in r.getMessage() and r.exc_info[0] is RuntimeError for r in reports) def a_document(size, dc_id=4): from telethon.tl import types return types.Document( id=1, access_hash=2, file_reference=b"r", date=datetime.now(timezone.utc), mime_type="video/x-matroska", size=size, dc_id=dc_id, attributes=[], ) async def test_a_parallel_download_is_counted_but_a_peek_at_the_start_of_a_file_is_not(tmp_path, monkeypatch): blob = blob_of(3 * 1024 * 1024 + 5) async def fake_open(client, dc_id, connections): assert dc_id == 4 return [FakeTransport(blob) for _ in range(2)] monkeypatch.setattr(ft, "open_transports", fake_open) await ft.fast_download(object(), a_document(len(blob)), tmp_path / "all.bin", connections=2) assert (tmp_path / "all.bin").read_bytes() == blob and ft.stats.fast_down == 1 await ft.fast_download(object(), a_document(len(blob)), tmp_path / "head.bin", connections=2, max_bytes=1024 * 1024) assert ft.stats.fast_down == 1 await ft.download(FakeClient(), Msg(a_document(len(blob))), tmp_path / "again.bin", fast=True) # the wrapper, end to end assert ft.stats == ft.TransferStats(fast_down=2) async def test_downloads_by_telethon_itself_are_counted_as_standard(tmp_path): await ft.download(FakeClient(), Msg(), tmp_path / "x.bin", fast=False) await ft.download(FakeClient(), Msg(document=None), tmp_path / "y.bin", fast=True) # no document to fetch in parallel await ft.download(FakeClient(), Msg(), tmp_path / "z.bin", fast=False, max_bytes=512 * 1024) # a peek is no transfer assert ft.stats == ft.TransferStats(standard_down=2) async def test_download_stock_progress_callback_can_cancel(tmp_path): cancel = asyncio.Event() cancel.set() with pytest.raises(Cancelled): await ft.download(FakeClient(), Msg(), tmp_path / "x", fast=False, cancel=cancel) async def test_download_partial_fallback_truncates(tmp_path, monkeypatch): async def broken(*a, **k): raise RuntimeError("nope") monkeypatch.setattr(ft, "fast_download", broken) dest = tmp_path / "head.bin" await ft.download(FakeClient(), Msg(), dest, fast=True, max_bytes=700 * 1024) assert dest.stat().st_size == 700 * 1024 async def test_upload_small_file_uses_stock_upload(tmp_path): src = tmp_path / "small.bin" src.write_bytes(b"x" * 1000) client = FakeClient() result = await ft.fast_upload(client, src, "My Name.mkv") assert result == ("stock-input-file", "My Name.mkv") assert client.uploaded == [(str(src), "My Name.mkv")] assert ft.stats == ft.TransferStats(small_up=1) # on purpose, so it is not counted as a failure or as parallel async def test_upload_wrapper_falls_back(tmp_path, monkeypatch): async def broken(*a, **k): raise RuntimeError("nope") monkeypatch.setattr(ft, "fast_upload", broken) src = tmp_path / "f.bin" src.write_bytes(b"x") client = FakeClient() assert await ft.upload(client, src, "n.mkv", fast=True) == ("stock-input-file", "n.mkv") assert (ft.stats.standard_up, ft.stats.fallbacks, ft.stats.last_fallback) == (1, 1, "upload: RuntimeError: nope") assert await ft.upload(client, src, "n.mkv", fast=False) == ("stock-input-file", "n.mkv") assert (ft.stats.standard_up, ft.stats.fallbacks) == (2, 1) # switched off is not a failure async def test_big_upload_builds_input_file_big(tmp_path, monkeypatch): size = ft.BIG_FILE_THRESHOLD + 12345 src = tmp_path / "big.bin" src.write_bytes(blob_of(size, seed=3)) store = {} async def fake_open(client, dc_id, connections): assert dc_id == 4 and connections == 3 return [FakeTransport(store=store, jitter=False) for _ in range(3)] monkeypatch.setattr(ft, "open_transports", fake_open) class C(FakeClient): class session: dc_id = 4 result = await ft.fast_upload(C(), src, "Big.mkv", connections=3) from telethon.tl.types import InputFileBig assert isinstance(result, InputFileBig) assert result.name == "Big.mkv" and result.parts == len(store) == -(-size // ft.UPLOAD_PART) assert result.id == next(iter(store.values()))[0] assert b"".join(store[i][2] for i in sorted(store)) == src.read_bytes() assert ft.stats == ft.TransferStats(fast_up=1) async def test_file_helper_without_positional_io(tmp_path): """Windows has no os.pread/os.pwrite: the lock + seek path must work too.""" f = ft._File(tmp_path / "w.bin", "w") f._positional = False f.pwrite(b"world", 5) f.pwrite(b"hello", 0) f.close() r = ft._File(tmp_path / "w.bin", "r") r._positional = False assert r.pread(10, 0) == b"helloworld" r.close() # --------------------------------------------------------------------------- # streamed copy: the download goes straight into the upload, nothing is saved # --------------------------------------------------------------------------- MIB = 1024 * 1024 @pytest.mark.parametrize("size", [1, ft.UPLOAD_PART - 1, ft.UPLOAD_PART, ft.UPLOAD_PART + 1, MIB - 1, MIB, MIB + 1, 3 * MIB // 2 + 7, 5 * MIB + 12345]) async def test_relay_copies_a_file_part_for_part(size): blob = blob_of(size, seed=11) store = {} down = [FakeTransport(blob) for _ in range(3)] up = [FakeTransport(store=store) for _ in range(3)] seen = [] n = await ft.relay_parts(down, up, "loc", size, 4321, progress=lambda c, t: seen.append((c, t))) expected = -(-size // ft.UPLOAD_PART) assert n == expected and sorted(store) == list(range(expected)) assert all(fid == 4321 and total == expected for fid, total, _d in store.values()) sizes = [len(store[i][2]) for i in range(expected)] assert all(s == ft.UPLOAD_PART for s in sizes[:-1]) # every part but the last is full size, whatever the download pieces were assert b"".join(store[i][2] for i in range(expected)) == blob assert seen[-1] == (size, size) and all(t == size for _c, t in seen) assert [c for c, _t in seen] == sorted(c for c, _t in seen) # progress only ever moves forward async def test_relay_fetches_every_part_once_and_both_sides_share_the_work(): size = 12 * MIB blob = blob_of(size) store = {} down = [FakeTransport(blob) for _ in range(4)] up = [FakeTransport(store=store) for _ in range(4)] await ft.relay_parts(down, up, "loc", size, 1) assert sorted(offset for t in down for offset, _limit in t.requests) == [i * MIB for i in range(12)] assert sum(t.calls for t in up) == 24 # two upload parts for every downloaded one assert sum(1 for t in down if t.calls) > 1 and sum(1 for t in up if t.calls) > 1 async def test_relay_with_far_more_connections_than_parts_still_finishes(): blob = blob_of(MIB + 100) # two downloaded parts, three upload parts store = {} down = [FakeTransport(blob) for _ in range(8)] up = [FakeTransport(store=store) for _ in range(8)] n = await asyncio.wait_for(ft.relay_parts(down, up, "loc", len(blob), 5), 5) assert n == 3 and b"".join(store[i][2] for i in sorted(store)) == blob async def test_relay_does_not_run_ahead_of_a_slow_upload(): parts, fetched, completed, ahead = 40, 0, 0, 0 blob = blob_of(parts * MIB) class Down(FakeTransport): async def get_file(self, location, offset, limit): nonlocal fetched, ahead data = await super().get_file(location, offset, limit) fetched += 1 ahead = max(ahead, fetched - completed) # downloaded parts that are not sent yet return data class Up(FakeTransport): async def save_part(self, file_id, part, total_parts, data): nonlocal completed await asyncio.sleep(0.005) await super().save_part(file_id, part, total_parts, data) if part % 2 == 1: # the second half of a downloaded part: that part is sent completed += 1 store = {} down, up = [Down(blob) for _ in range(4)], [Up(store=store) for _ in range(2)] await ft.relay_parts(down, up, "loc", parts * MIB, 1) held_back = max(ft.RELAY_BUFFER_MIN, 2 * len(up)) # the buffer, plus a part in each hand assert ahead <= held_back + len(down) + len(up) and ahead < parts // 2 # memory stays small however big the file is assert b"".join(store[i][2] for i in sorted(store)) == blob async def test_relay_retries_failures_on_both_sides(): size = 3 * MIB + 7 blob = blob_of(size) store = {} down = [FakeTransport(blob, fail_first=2), FakeTransport(blob, fail_first=3, short_once=True)] up = [FakeTransport(store=store, fail_first=2), FakeTransport(store=store)] await ft.relay_parts(down, up, "loc", size, 1) assert b"".join(store[i][2] for i in sorted(store)) == blob async def test_relay_gives_up_when_a_connection_keeps_failing(): blob = blob_of(2 * MIB) broken = FakeTransport(store={}, fail_first=10_000) with pytest.raises(ConnectionError): await asyncio.wait_for(ft.relay_parts([FakeTransport(blob)], [broken], "loc", len(blob), 1), 10) assert broken.calls == ft.MAX_RETRIES + 1 @pytest.mark.parametrize("broken_side", ["down", "up"]) async def test_relay_stops_both_sides_when_one_fails_for_good(broken_side): class FileReferenceExpiredError(Exception): # retrying cannot help pass class LateFatalUp(FakeTransport): async def save_part(self, file_id, part, total_parts, data): await asyncio.sleep(0.05) # long enough for the downloads to fill the buffer and wait raise FileReferenceExpiredError("expired") blob = blob_of(40 * MIB) if broken_side == "down": down, up = [FakeTransport(blob, fatal=FileReferenceExpiredError("expired"))], [FakeTransport(store={}) for _ in range(2)] else: down, up = [FakeTransport(blob) for _ in range(3)], [LateFatalUp() for _ in range(2)] with pytest.raises(FileReferenceExpiredError): # and does not hang on a full or an empty buffer await asyncio.wait_for(ft.relay_parts(down, up, "loc", len(blob), 1), 5) assert sum(t.calls for t in down) < 40 # the download did not carry on to the end async def test_relay_cancel_stops_both_sides(): cancel = asyncio.Event() class SlowDown(FakeTransport): async def get_file(self, location, offset, limit): await asyncio.sleep(0.005) if offset >= 3 * MIB: cancel.set() return b"\0" * limit up = [FakeTransport(store={}) for _ in range(3)] with pytest.raises(Cancelled): await asyncio.wait_for(ft.relay_parts([SlowDown() for _ in range(3)], up, "loc", 64 * MIB, 1, cancel=cancel), 5) assert sum(t.calls for t in up) < 2 * 64 cancel.set() # cancelled before it starts untouched = FakeTransport(store={}) with pytest.raises(Cancelled): await ft.relay_parts([FakeTransport(b"x" * MIB)], [untouched], "loc", MIB, 1, cancel=cancel) assert untouched.calls == 0 async def test_relay_refuses_what_it_cannot_do(): ok = [FakeTransport(b"x" * 10)] with pytest.raises(ValueError): # an empty file has no parts await ft.relay_parts(ok, ok, "loc", 0, 1) with pytest.raises(ValueError): # no connection to fetch with, or none to send with await ft.relay_parts([], ok, "loc", 10, 1) with pytest.raises(ValueError): await ft.relay_parts(ok, [], "loc", 10, 1) with pytest.raises(ValueError): # Telegram wants equal parts: a download part must be cut into whole ones await ft.relay_parts(ok, ok, "loc", 10, 1, part_size=1000, send_size=512) class RelayClient: class session: dc_id = 2 async def test_relay_fetches_from_the_files_dc_sends_to_ours_and_returns_an_input_file_big(monkeypatch): from telethon.tl.types import InputFileBig size = ft.BIG_FILE_THRESHOLD + 12345 blob = blob_of(size, seed=5) store, opened = {}, [] async def fake_open(client, dc_id, connections): group = [FakeTransport(blob, store=store, jitter=False) for _ in range(connections)] opened.append((dc_id, group)) return group monkeypatch.setattr(ft, "open_transports", fake_open) seen = [] result = await ft.relay(RelayClient(), a_document(size, dc_id=4), "Big Name.mkv", connections=3, progress=lambda c, t: seen.append(c)) assert sorted(dc for dc, _group in opened) == [2, 4] and all(len(g) == 3 for _dc, g in opened) assert isinstance(result, InputFileBig) and result.name == "Big Name.mkv" assert result.parts == len(store) == -(-size // ft.UPLOAD_PART) and result.id == next(iter(store.values()))[0] assert b"".join(store[i][2] for i in sorted(store)) == blob assert all(t.closed for _dc, group in opened for t in group) and seen[-1] == size assert ft.stats == ft.TransferStats() # counting what was streamed is up to the caller, once the file is sent async def test_relay_closes_its_connections_when_it_fails_and_when_it_cannot_open_them(monkeypatch): opened = [] async def failing_open(client, dc_id, connections): group = [FakeTransport(b"", fail_first=10_000) for _ in range(connections)] opened.append(group) return group monkeypatch.setattr(ft, "open_transports", failing_open) with pytest.raises(ConnectionError): await ft.relay(RelayClient(), a_document(ft.BIG_FILE_THRESHOLD + 1, dc_id=4), "x.mkv", connections=2) assert len(opened) == 2 and all(t.closed for group in opened for t in group) opened.clear() async def half_open(client, dc_id, connections): if dc_id == 4: raise OSError("no route to DC 4") group = [FakeTransport() for _ in range(connections)] opened.append(group) return group monkeypatch.setattr(ft, "open_transports", half_open) with pytest.raises(OSError): await ft.relay(RelayClient(), a_document(ft.BIG_FILE_THRESHOLD + 1, dc_id=4), "x.mkv", connections=2) assert len(opened) == 1 and all(t.closed for t in opened[0]) # what did open is not left connected def test_file_part_errors_are_told_apart_from_the_rest(): from telethon import errors assert ft.is_file_part_error(errors.FilePartsInvalidError(request=None)) assert ft.is_file_part_error(errors.FilePartMissingError(request=None, capture=3)) assert not ft.is_file_part_error(errors.MediaCaptionTooLongError(request=None)) assert not ft.is_file_part_error(RuntimeError("FilePart")) def test_a_failed_stream_is_remembered_apart_from_the_parallel_failures(): stats = ft.TransferStats() stats.stream_failed(ValueError("boom")) assert (stats.stream_fallbacks, stats.last_stream_fallback, stats.fallbacks, stats.last_fallback) == (1, "ValueError: boom", 0, "")