Spaces:
Running
Running
| import asyncio | |
| import random | |
| from datetime import datetime, timezone | |
| import pytest | |
| from bot import fasttransfer as ft | |
| from bot.errors import Cancelled | |
| def fast_retries(monkeypatch): | |
| monkeypatch.setattr(ft, "RETRY_BASE_DELAY", 0.0) | |
| 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) | |
| 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] | |
| 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 | |
| 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 | |
| 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, "") | |