ArchiveBox/archivebox/tests/test_snapshot_service.py

import json
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import pytest
from django.db import close_old_connections

from archivebox.core.models import ArchiveResult, Snapshot
from archivebox.tests.conftest import run_archivebox_cmd
from archivebox.tests.test_orm_helpers import use_archivebox_db
from .conftest import (
    cli_env,
    get_free_port,
    init_archive,
)

pytestmark = pytest.mark.django_db(transaction=True)


def _snapshot_state(cwd: Path, url: str) -> dict[str, object]:
    with use_archivebox_db(cwd):
        snapshot = Snapshot.objects.select_related("crawl", "crawl__created_by").get(url=url)
        snapshot_dir = Path(snapshot.output_dir)
        crawl_dir = Path(snapshot.crawl.output_dir)
        crawl_link = crawl_dir / "snapshots" / Snapshot.extract_domain_from_url(snapshot.url) / str(snapshot.id)
        results = list(
            ArchiveResult.objects.filter(snapshot=snapshot)
            .order_by("plugin", "hook_name")
            .values("plugin", "hook_name", "status", "output_files", "output_size"),
        )
        return {
            "id": str(snapshot.id),
            "crawl_id": str(snapshot.crawl_id),
            "status": snapshot.status,
            "retry_at": snapshot.retry_at,
            "downloaded_at": snapshot.downloaded_at,
            "output_size": snapshot.output_size,
            "snapshot_dir": snapshot_dir,
            "crawl_dir": crawl_dir,
            "crawl_link": crawl_link,
            "results": results,
        }


def test_snapshot_completion_preserves_retry_scheduled_during_active_run(tmp_path, admin_user):
    from datetime import timedelta

    from django.utils import timezone

    from archivebox.crawls.models import Crawl
    from archivebox.services.snapshot_service import finalize_completed_snapshot

    crawl = Crawl.objects.create(urls="https://example.com/retry-race", created_by=admin_user)
    owned_retry_at = timezone.now() + timedelta(minutes=10)
    snapshot = Snapshot.objects.create(
        url="https://example.com/retry-race",
        crawl=crawl,
        status=Snapshot.StatusChoices.STARTED,
        retry_at=owned_retry_at,
        config={"RETRY_PLUGINS": ["title"]},
    )

    snapshot.schedule_plugin_run(["title"])
    finalize_completed_snapshot(
        str(snapshot.id),
        owned_retry_at=owned_retry_at,
        was_sealed=False,
        consumed_retry_plugins=["title"],
        output_dir=tmp_path,
    )

    snapshot.refresh_from_db()
    assert snapshot.status == Snapshot.StatusChoices.QUEUED
    assert snapshot.retry_at is not None
    assert snapshot.retry_at < owned_retry_at
    assert snapshot.config["RETRY_PLUGINS"] == ["title"]


def test_concurrent_plugin_scheduling_durably_merges_every_request(admin_user):
    from archivebox.crawls.models import Crawl
    from django.db import connection

    if connection.vendor != "sqlite":
        pytest.skip("exercises SQLite concurrent-writer scheduling")

    crawl = Crawl.objects.create(urls="https://example.com/plugin-race", created_by=admin_user)
    snapshot = Snapshot.objects.create(url="https://example.com/plugin-race", crawl=crawl)
    plugins = ["title", "wget", "screenshot", "pdf"]

    def schedule(plugin: str) -> bool:
        close_old_connections()
        try:
            return Snapshot.objects.get(pk=snapshot.pk).schedule_plugin_run([plugin])
        finally:
            close_old_connections()

    with ThreadPoolExecutor(max_workers=len(plugins)) as pool:
        results = list(pool.map(schedule, plugins))

    snapshot.refresh_from_db()
    assert results == [True] * len(plugins)
    assert snapshot.config["RETRY_PLUGINS"] == sorted(plugins)


def test_snapshot_keyset_iterator_reads_more_than_eight_pages(admin_user):
    from archivebox.crawls.models import Crawl

    crawl = Crawl.objects.create(urls="https://example.com/pages", created_by=admin_user)
    snapshots = [Snapshot.objects.create(url=f"https://example.com/pages/{idx}", crawl=crawl) for idx in range(10)]

    yielded_ids = [snapshot.id for snapshot in crawl.snapshot_set.order_by("id").paged_iterator(chunk_size=1)]

    assert yielded_ids == sorted(snapshot.id for snapshot in snapshots)


def test_snapshot_merge_consolidates_only_exact_hook_identity(admin_user):
    from archivebox.crawls.models import Crawl

    keeper_crawl = Crawl.objects.create(urls="https://example.com", created_by=admin_user)
    duplicate_crawl = Crawl.objects.create(urls="https://example.com", created_by=admin_user)
    keeper = Snapshot.objects.create(url="https://example.com", crawl=keeper_crawl)
    duplicate = Snapshot.objects.create(url="https://example.com", crawl=duplicate_crawl)
    ArchiveResult.objects.create(
        snapshot=keeper,
        plugin="archivewebpage",
        hook_name="on_Snapshot__65_archivewebpage_stop",
        status=ArchiveResult.StatusChoices.NORESULTS,
        output_files={"older.wacz": {"size": 5}},
        output_size=5,
    )
    ArchiveResult.objects.create(
        snapshot=duplicate,
        plugin="archivewebpage",
        hook_name="on_Snapshot__65_archivewebpage_stop",
        status=ArchiveResult.StatusChoices.SUCCEEDED,
        output_str="archivewebpage.wacz",
        output_files={"archivewebpage.wacz": {"size": 7}},
        output_size=7,
    )
    ArchiveResult.objects.create(
        snapshot=duplicate,
        plugin="archivewebpage",
        hook_name="on_Snapshot__16_archivewebpage_start",
        status=ArchiveResult.StatusChoices.SUCCEEDED,
        output_str="recording started",
    )

    Snapshot._merge_snapshots([keeper, duplicate])

    assert not Snapshot.objects.filter(pk=duplicate.pk).exists()
    results = ArchiveResult.objects.filter(snapshot=keeper, plugin="archivewebpage")
    assert results.count() == 2
    stop_result = results.get(hook_name="on_Snapshot__65_archivewebpage_stop")
    assert stop_result.status == ArchiveResult.StatusChoices.SUCCEEDED
    assert stop_result.output_str == "archivewebpage.wacz"
    assert set(stop_result.output_files) == {"older.wacz", "archivewebpage.wacz"}
    assert results.filter(hook_name="on_Snapshot__16_archivewebpage_start").exists()


@pytest.mark.timeout(180)
def test_snapshot_service_cli_add_seals_snapshot_and_writes_indexes(tmp_path, recursive_test_site):
    init_archive(tmp_path)

    port = get_free_port()
    env = cli_env(port=port, server=True, PLUGINS="wget", SAVE_WGET="True")
    _cmd_result = run_archivebox_cmd(
        ["add", "--depth=0", "--plugins=wget", recursive_test_site["root_url"]],
        cwd=tmp_path,
        env=env,
        timeout=180,
    )
    stdout, stderr, code = _cmd_result.stdout, _cmd_result.stderr, _cmd_result.returncode
    assert code == 0, f"archivebox add failed with code {code}\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}"

    state = _snapshot_state(tmp_path, recursive_test_site["root_url"])
    snapshot_dir = state["snapshot_dir"]
    crawl_link = state["crawl_link"]
    assert isinstance(snapshot_dir, Path)
    assert isinstance(crawl_link, Path)

    assert state["status"] == Snapshot.StatusChoices.SEALED
    assert state["retry_at"] is None
    assert state["downloaded_at"] is not None
    assert snapshot_dir.is_dir()
    assert crawl_link.is_symlink()
    assert crawl_link.resolve() == snapshot_dir.resolve()

    index_jsonl = snapshot_dir / "index.jsonl"
    assert index_jsonl.is_file()

    records = [json.loads(line) for line in index_jsonl.read_text(encoding="utf-8").splitlines() if line.strip()]
    assert records[0]["type"] == "Snapshot"
    assert records[0]["id"] == state["id"]
    assert any(record.get("type") == "ArchiveResult" and record.get("plugin") == "wget" for record in records)

    wget_files = [path for path in (snapshot_dir / "wget").rglob("*") if path.is_file()]
    assert wget_files
    assert any("Root" in path.read_text(encoding="utf-8", errors="ignore") for path in wget_files if path.suffix in (".html", ".txt"))
    assert any(result["plugin"] == "wget" and result["status"] == ArchiveResult.StatusChoices.SUCCEEDED for result in state["results"])