tests.py

"""Fetch many URLs concurrently with the async client.

Usage:
    export FASTSAVER_API_KEY="fs_sk_..."
    python examples/async_batch.py <url1> <url2> ...

Without arguments a small built-in list of URLs is used. Each URL is fetched
concurrently with ``asyncio.gather(..., return_exceptions=True)`` so a single
failure (or rate limit) never aborts the whole batch.
"""

from __future__ import annotations

import asyncio
import sys

from fastsaver import AsyncFastSaver, FastSaverError, FetchError, MediaResult, RateLimitError

DEFAULT_URLS = [
    "https://www.instagram.com/reel/DRsmm9UjKfH/",
    "https://www.tiktok.com/@scout2015/video/6718335390845095173",
    "https://pin.it/1a2b3c4d5",
    "https://x.com/jack/status/20",
]


def describe(url: str, outcome: MediaResult | BaseException) -> str:
    """Turn a gather() outcome (a MediaResult or an exception) into a short report."""
    if isinstance(outcome, MediaResult):
        urls = outcome.download_urls
        return f"OK    {url}\n      {outcome.type} from {outcome.source}, {len(urls)} file(s)"
    if isinstance(outcome, RateLimitError):
        hint = f" (retry after {outcome.retry_after}s)" if outcome.retry_after else ""
        return f"LIMIT {url}\n      rate limit exceeded{hint}"
    if isinstance(outcome, FetchError):
        # The service could not resolve the media (private post, deleted, ...).
        return f"FAIL  {url}\n      {outcome.code}"
    if isinstance(outcome, FastSaverError):
        return f"ERROR {url}\n      {outcome}"
    # Anything else is an unexpected exception; re-raise so it is not swallowed.
    raise outcome


async def main(urls: list[str]) -> int:
    async with AsyncFastSaver() as client:
        outcomes = await asyncio.gather(
            *(client.fetch(url) for url in urls),
            return_exceptions=True,
        )

    failures = 0
    for url, outcome in zip(urls, outcomes):
        print(describe(url, outcome))
        if isinstance(outcome, BaseException):
            failures += 1

    print(f"\n{len(urls) - failures}/{len(urls)} succeeded")
    return 1 if failures else 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main(sys.argv[1:] or DEFAULT_URLS)))