140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
"""End-to-end test: the real downloader and HTTP layer against a local server.
|
|
|
|
Only the Telegram Bot API itself is faked; everything else — planning, the
|
|
concurrency limit, retries, atomic writes and metadata — is the real code path.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
|
|
import aiohttp
|
|
import pytest
|
|
import pytest_asyncio
|
|
from aiohttp import web
|
|
|
|
from sticker_downloader.config import METADATA_FILE_NAME, DownloadConfig
|
|
from sticker_downloader.downloader import StickerDownloader
|
|
from sticker_downloader.fetcher import AiohttpFetcher
|
|
from tests.conftest import (
|
|
FakeFile,
|
|
FakeStickerSet,
|
|
FakeThumbnail,
|
|
make_sticker,
|
|
no_sleep,
|
|
)
|
|
|
|
#: file_id -> (extension, body). "flaky1" fails once before succeeding.
|
|
FILES = {
|
|
"file1": ("webp", b"one"),
|
|
"thumb1": ("png", b"one-thumb"),
|
|
"file2": ("webp", b"two"),
|
|
"thumb2": ("png", b"two-thumb"),
|
|
"packthumb": ("webp", b"cover"),
|
|
}
|
|
|
|
|
|
class _Server:
|
|
def __init__(self) -> None:
|
|
self.base = ""
|
|
self.hits: list[str] = []
|
|
self.transient_failures = {"file2"}
|
|
|
|
async def handle(self, request: web.Request) -> web.Response:
|
|
file_id = request.match_info["file_id"]
|
|
self.hits.append(file_id)
|
|
if file_id in self.transient_failures:
|
|
# Fail the first attempt only, so the retry path runs for real.
|
|
self.transient_failures.discard(file_id)
|
|
return web.Response(status=503, text="try again")
|
|
extension, body = FILES[file_id]
|
|
return web.Response(body=body, content_type=f"image/{extension}")
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def server() -> AsyncIterator[_Server]:
|
|
state = _Server()
|
|
app = web.Application()
|
|
app.router.add_get("/{file_id}.{extension}", state.handle)
|
|
runner = web.AppRunner(app)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, "127.0.0.1", 0)
|
|
await site.start()
|
|
state.base = f"http://127.0.0.1:{runner.addresses[0][1]}"
|
|
try:
|
|
yield state
|
|
finally:
|
|
await runner.cleanup()
|
|
|
|
|
|
class ServingBot:
|
|
"""A Bot stand-in whose file paths point at the local test server."""
|
|
|
|
def __init__(self, base_url: str, sets: dict[str, FakeStickerSet]) -> None:
|
|
self._base_url = base_url
|
|
self._sets = sets
|
|
|
|
async def get_sticker_set(self, name: str) -> FakeStickerSet:
|
|
return self._sets[name]
|
|
|
|
async def get_file(self, file_id: str) -> FakeFile:
|
|
extension, _body = FILES[file_id]
|
|
return FakeFile(file_path=f"{self._base_url}/{file_id}.{extension}")
|
|
|
|
|
|
@pytest.fixture
|
|
def pack_of_two() -> FakeStickerSet:
|
|
return FakeStickerSet(
|
|
name="RealPack",
|
|
title="Real Pack",
|
|
stickers=[make_sticker(1), make_sticker(2)],
|
|
thumbnail=FakeThumbnail(file_id="packthumb"),
|
|
)
|
|
|
|
|
|
async def test_full_download_over_http(server, pack_of_two, tmp_path):
|
|
config = DownloadConfig(output_dir=tmp_path / "out", concurrency=2)
|
|
async with aiohttp.ClientSession() as session:
|
|
downloader = StickerDownloader(
|
|
bot=ServingBot(server.base, {"RealPack": pack_of_two}),
|
|
fetcher=AiohttpFetcher(session),
|
|
config=config,
|
|
sleep=no_sleep,
|
|
)
|
|
result = await downloader.download_pack("https://t.me/addstickers/RealPack")
|
|
|
|
assert result.ok, result.error
|
|
assert result.downloaded == 5
|
|
|
|
directory = config.output_dir / "RealPack"
|
|
assert (directory / "001.webp").read_bytes() == b"one"
|
|
assert (directory / "001_thumb.png").read_bytes() == b"one-thumb"
|
|
assert (directory / "002.webp").read_bytes() == b"two"
|
|
assert (directory / "_pack_thumbnail.webp").read_bytes() == b"cover"
|
|
assert list(directory.glob("*.part")) == []
|
|
|
|
# The 503 on file2 was retried, so that file was requested twice.
|
|
assert server.hits.count("file2") == 2
|
|
|
|
metadata = json.loads((directory / METADATA_FILE_NAME).read_text(encoding="utf-8"))
|
|
assert metadata["sticker_count"] == 2
|
|
assert metadata["stickers"][1]["files"] == ["002.webp", "002_thumb.png"]
|
|
|
|
|
|
async def test_second_run_skips_everything(server, pack_of_two, tmp_path):
|
|
config = DownloadConfig(output_dir=tmp_path / "out")
|
|
bot = ServingBot(server.base, {"RealPack": pack_of_two})
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
downloader = StickerDownloader(
|
|
bot=bot, fetcher=AiohttpFetcher(session), config=config, sleep=no_sleep
|
|
)
|
|
await downloader.download_pack("RealPack")
|
|
hits_after_first = len(server.hits)
|
|
second = await downloader.download_pack("RealPack")
|
|
|
|
assert second.downloaded == 0
|
|
assert second.skipped == 5
|
|
assert len(server.hits) == hits_after_first # no new requests at all
|