271 lines
9.0 KiB
Python
271 lines
9.0 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import replace
|
|
|
|
import pytest
|
|
from telegram.error import TimedOut
|
|
|
|
from sticker_downloader.config import METADATA_FILE_NAME, DownloadConfig
|
|
from sticker_downloader.downloader import (
|
|
StickerDownloader,
|
|
primary_file_type,
|
|
sanitize_filename_part,
|
|
)
|
|
from sticker_downloader.errors import FetchError
|
|
from sticker_downloader.results import FileStatus
|
|
from tests.conftest import (
|
|
FakeBot,
|
|
FakeFetcher,
|
|
FakeSticker,
|
|
FakeStickerSet,
|
|
FakeThumbnail,
|
|
make_sticker,
|
|
no_sleep,
|
|
)
|
|
|
|
|
|
def build(bot, fetcher, config, **kwargs) -> StickerDownloader:
|
|
return StickerDownloader(
|
|
bot=bot, fetcher=fetcher, config=config, sleep=no_sleep, **kwargs
|
|
)
|
|
|
|
|
|
async def test_downloads_stickers_thumbnails_and_pack_cover(bot, fetcher, config):
|
|
result = await build(bot, fetcher, config).download_pack(
|
|
"https://t.me/addstickers/TestPack"
|
|
)
|
|
|
|
assert result.ok
|
|
assert result.name == "TestPack"
|
|
assert result.title == "Test Pack"
|
|
assert result.total_stickers == 2
|
|
|
|
directory = config.output_dir / "TestPack"
|
|
written = sorted(p.name for p in directory.iterdir())
|
|
assert written == [
|
|
"001.webp",
|
|
"001_thumb.png",
|
|
"002.webp",
|
|
"002_thumb.png",
|
|
"_pack_thumbnail.webp",
|
|
METADATA_FILE_NAME,
|
|
]
|
|
assert (directory / "001.webp").read_bytes() == b"sticker-bytes"
|
|
assert result.downloaded == 5
|
|
assert result.bytes_written == 5 * len(b"sticker-bytes")
|
|
|
|
|
|
async def test_file_type_filter_limits_what_is_written(bot, fetcher, config):
|
|
config = replace(config, file_types=frozenset({"webp"}), pack_thumbnail=False)
|
|
result = await build(bot, fetcher, config).download_pack("TestPack")
|
|
|
|
directory = config.output_dir / "TestPack"
|
|
assert sorted(p.name for p in directory.iterdir()) == [
|
|
"001.webp",
|
|
"002.webp",
|
|
METADATA_FILE_NAME,
|
|
]
|
|
assert result.downloaded == 2
|
|
# No thumbnails were even looked up.
|
|
assert bot.get_file_calls == ["file1", "file2"]
|
|
|
|
|
|
async def test_extension_follows_sticker_kind(config, fetcher):
|
|
animated = FakeSticker(file_id="anim", is_animated=True)
|
|
video = FakeSticker(file_id="vid", is_video=True)
|
|
pack = FakeStickerSet(name="Mixed", stickers=[animated, video])
|
|
bot = FakeBot({"Mixed": pack}, extensions={"anim": "tgs", "vid": "webm"})
|
|
|
|
await build(bot, fetcher, config).download_pack("Mixed")
|
|
|
|
directory = config.output_dir / "Mixed"
|
|
assert (directory / "001.tgs").is_file()
|
|
assert (directory / "002.webm").is_file()
|
|
|
|
|
|
async def test_pack_thumbnail_extension_comes_from_telegram(config, fetcher):
|
|
pack = FakeStickerSet(
|
|
name="Vid", stickers=[make_sticker(1)], thumbnail=FakeThumbnail("cover")
|
|
)
|
|
bot = FakeBot({"Vid": pack}, extensions={"cover": "webm"})
|
|
|
|
await build(bot, fetcher, config).download_pack("Vid")
|
|
|
|
assert (config.output_dir / "Vid" / "_pack_thumbnail.webm").is_file()
|
|
|
|
|
|
async def test_existing_files_are_skipped_without_api_calls(bot, fetcher, config):
|
|
directory = config.output_dir / "TestPack"
|
|
directory.mkdir(parents=True)
|
|
(directory / "001.webp").write_bytes(b"old")
|
|
|
|
result = await build(bot, fetcher, config).download_pack("TestPack")
|
|
|
|
assert result.skipped == 1
|
|
assert result.downloaded == 4
|
|
assert (directory / "001.webp").read_bytes() == b"old"
|
|
assert "file1" not in bot.get_file_calls
|
|
|
|
|
|
async def test_overwrite_replaces_existing_files(bot, fetcher, config):
|
|
directory = config.output_dir / "TestPack"
|
|
directory.mkdir(parents=True)
|
|
(directory / "001.webp").write_bytes(b"old")
|
|
|
|
config = replace(config, overwrite=True)
|
|
result = await build(bot, fetcher, config).download_pack("TestPack")
|
|
|
|
assert result.skipped == 0
|
|
assert (directory / "001.webp").read_bytes() == b"sticker-bytes"
|
|
|
|
|
|
async def test_dry_run_writes_nothing(bot, fetcher, config):
|
|
config = replace(config, dry_run=True)
|
|
result = await build(bot, fetcher, config).download_pack("TestPack")
|
|
|
|
assert result.planned == 5
|
|
assert result.downloaded == 0
|
|
assert not config.output_dir.exists()
|
|
assert fetcher.urls == []
|
|
assert bot.get_file_calls == []
|
|
|
|
|
|
async def test_invalid_reference_never_touches_the_network(bot, fetcher, config):
|
|
result = await build(bot, fetcher, config).download_pack("https://example.com/nope")
|
|
|
|
assert not result.ok
|
|
assert result.error is not None
|
|
assert bot.get_sticker_set_calls == []
|
|
|
|
|
|
async def test_unknown_pack_reports_a_readable_error(bot, fetcher, config):
|
|
result = await build(bot, fetcher, config).download_pack("NoSuchPack")
|
|
|
|
assert not result.ok
|
|
assert result.error == "sticker pack not found"
|
|
assert not config.output_dir.exists()
|
|
|
|
|
|
async def test_one_bad_file_does_not_sink_the_pack(bot, config):
|
|
fetcher = FakeFetcher(failures={"file1": FetchError("file1", 404, "Not Found")})
|
|
result = await build(bot, fetcher, config).download_pack("TestPack")
|
|
|
|
assert not result.ok
|
|
assert result.failed == 1
|
|
assert result.downloaded == 4
|
|
failure = next(o for o in result.outcomes if o.status is FileStatus.FAILED)
|
|
assert failure.path.name == "001.webp"
|
|
assert "404" in (failure.error or "")
|
|
|
|
|
|
async def test_transient_failures_are_retried(bot, config):
|
|
calls = {"n": 0}
|
|
real_fetch = FakeFetcher().fetch
|
|
|
|
class FlakyFetcher:
|
|
async def fetch(self, url: str) -> bytes:
|
|
calls["n"] += 1
|
|
if calls["n"] == 1:
|
|
raise TimedOut
|
|
return await real_fetch(url)
|
|
|
|
result = await build(bot, FlakyFetcher(), config).download_pack("TestPack")
|
|
|
|
assert result.ok
|
|
assert result.downloaded == 5
|
|
assert calls["n"] == 6 # five files plus the one retry
|
|
|
|
|
|
async def test_no_partial_files_are_left_behind(bot, config):
|
|
fetcher = FakeFetcher(failures={"file2": FetchError("file2", 500)})
|
|
config = replace(config, retries=1)
|
|
await build(bot, fetcher, config).download_pack("TestPack")
|
|
|
|
directory = config.output_dir / "TestPack"
|
|
assert [p.name for p in directory.glob("*.part")] == []
|
|
|
|
|
|
async def test_concurrency_is_bounded(config, fetcher):
|
|
stickers = [make_sticker(i) for i in range(1, 11)]
|
|
bot = FakeBot({"Big": FakeStickerSet(name="Big", stickers=stickers)})
|
|
config = replace(config, concurrency=3)
|
|
|
|
await build(bot, fetcher, config).download_pack("Big")
|
|
|
|
assert fetcher.max_concurrent <= 3
|
|
assert len(fetcher.urls) == 20
|
|
|
|
|
|
async def test_emoji_names_are_used_when_requested(config, fetcher):
|
|
pack = FakeStickerSet(name="Emo", stickers=[FakeSticker(file_id="a", emoji="🦊")])
|
|
bot = FakeBot({"Emo": pack})
|
|
config = replace(config, emoji_names=True, pack_thumbnail=False)
|
|
|
|
await build(bot, fetcher, config).download_pack("Emo")
|
|
|
|
assert (config.output_dir / "Emo" / "001_🦊.webp").is_file()
|
|
|
|
|
|
async def test_metadata_describes_the_pack(bot, fetcher, config):
|
|
await build(bot, fetcher, config).download_pack("https://t.me/addstickers/TestPack")
|
|
|
|
payload = json.loads(
|
|
(config.output_dir / "TestPack" / METADATA_FILE_NAME).read_text(encoding="utf-8")
|
|
)
|
|
assert payload["pack"]["name"] == "TestPack"
|
|
assert payload["pack"]["title"] == "Test Pack"
|
|
assert payload["pack"]["reference"] == "https://t.me/addstickers/TestPack"
|
|
assert payload["sticker_count"] == 2
|
|
assert payload["requested_file_types"] == ["webp", "tgs", "webm", "png"]
|
|
first = payload["stickers"][0]
|
|
assert first["index"] == 1
|
|
assert first["emoji"] == "😀"
|
|
assert first["files"] == ["001.webp", "001_thumb.png"]
|
|
|
|
|
|
async def test_metadata_can_be_disabled(bot, fetcher, config):
|
|
config = replace(config, write_metadata=False)
|
|
await build(bot, fetcher, config).download_pack("TestPack")
|
|
assert not (config.output_dir / "TestPack" / METADATA_FILE_NAME).exists()
|
|
|
|
|
|
async def test_download_all_reports_every_pack(fetcher, config):
|
|
packs = {
|
|
"One": FakeStickerSet(name="One", stickers=[make_sticker(1)]),
|
|
"Two": FakeStickerSet(name="Two", stickers=[make_sticker(1)]),
|
|
}
|
|
bot = FakeBot(packs)
|
|
|
|
results = await build(bot, fetcher, config).download_all(["One", "Two", "Missing"])
|
|
|
|
assert [r.reference for r in results] == ["One", "Two", "Missing"]
|
|
assert [r.ok for r in results] == [True, True, False]
|
|
|
|
|
|
def test_primary_file_type_prefers_animated_over_video():
|
|
assert primary_file_type(FakeSticker(file_id="a")) == "webp"
|
|
assert primary_file_type(FakeSticker(file_id="a", is_animated=True)) == "tgs"
|
|
assert primary_file_type(FakeSticker(file_id="a", is_video=True)) == "webm"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("raw", "expected"),
|
|
[
|
|
("🦊", "🦊"),
|
|
("a/b", "ab"),
|
|
("..", ""),
|
|
('x<>:"|?*y', "xy"),
|
|
("x" * 50, "x" * 32),
|
|
],
|
|
)
|
|
def test_sanitize_filename_part(raw: str, expected: str):
|
|
assert sanitize_filename_part(raw) == expected
|
|
|
|
|
|
async def test_output_directory_is_created_lazily(bot, fetcher, tmp_path):
|
|
nested = tmp_path / "a" / "b" / "c"
|
|
config = DownloadConfig(output_dir=nested)
|
|
await build(bot, fetcher, config).download_pack("TestPack")
|
|
assert (nested / "TestPack" / "001.webp").is_file()
|