ArchiveBox/archivebox/tests/test_binary_service.py

import json
import os
import shutil
import uuid
from pathlib import Path

import pytest

from archivebox.machine.models import Binary, Machine, Process
from archivebox.tests.conftest import parse_jsonl_output, run_archivebox_cmd
from archivebox.tests.test_orm_helpers import use_archivebox_db

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


def _link_real_binary(bin_dir: Path, name: str, *, source: str | None = None) -> Path:
    bin_dir.mkdir(parents=True, exist_ok=True)
    source_path = shutil.which(source or name)
    assert source_path, f"{source or name} must be installed for this integration test"
    link = bin_dir / name
    link.unlink(missing_ok=True)
    link.symlink_to(source_path)
    return link


def _binary_request(name: str, *, binproviders: str = "env") -> str:
    return json.dumps({"type": "BinaryRequest", "name": name, "binproviders": binproviders}) + "\n"


def _runtime_env(data_dir: Path, bin_dir: Path) -> dict[str, str]:
    return {
        "LIB_DIR": str(data_dir / "lib"),
        "LIB_BIN_DIR": str(data_dir / "lib" / "bin"),
        "ABXPKG_LIB_DIR": str(data_dir / "lib"),
        "PATH": os.pathsep.join([str(bin_dir), "/usr/bin", "/bin", "/usr/sbin", "/sbin"]),
    }


def test_binary_request_installs_env_binary_and_recovers_stale_cache(initialized_archive, tmp_path):
    name = f"abx-e2e-rg-{uuid.uuid4().hex[:8]}"
    bootstrap_bin_dir = tmp_path / "realbin"
    provider_bin_dir = initialized_archive / "lib" / "env" / "bin"
    _link_real_binary(bootstrap_bin_dir, "uv")
    _link_real_binary(provider_bin_dir, name, source="rg")

    stdout, stderr, returncode = run_archivebox_cmd(
        ["run"],
        data_dir=initialized_archive,
        stdin=_binary_request(name),
        timeout=120,
        env=_runtime_env(initialized_archive, bootstrap_bin_dir),
    )

    assert returncode == 0, stderr
    output_records = parse_jsonl_output(stdout)
    assert any(record["type"] == "BinaryRequest" and record["name"] == name for record in output_records)

    with use_archivebox_db(initialized_archive):
        binary = Binary.objects.get(name=name)
        machine_id = str(binary.machine_id)
        first_binary_id = str(binary.id)
        first_abspath = Path(binary.abspath)
        binary_processes = list(Process.objects.filter(process_type=Process.TypeChoices.BINARY).order_by("created_at"))

    assert binary.status == Binary.StatusChoices.INSTALLED
    assert binary.version
    assert binary.binprovider == "env"
    assert binary.binproviders == "env"
    assert first_abspath.exists()
    assert first_abspath == provider_bin_dir / name
    assert first_abspath.resolve() == Path(shutil.which("rg") or "").resolve()
    assert first_abspath.is_relative_to(initialized_archive / "lib")
    assert (initialized_archive / "lib" / "env" / "bin" / name).exists()
    assert (initialized_archive / "lib" / "bin" / "rg").exists()
    assert (initialized_archive / "machines" / machine_id / "binaries" / name / "index.jsonl").exists()
    assert binary_processes
    assert binary_processes[-1].status == Process.StatusChoices.EXITED
    assert binary_processes[-1].exit_code == 0
    assert any(f"--name={name}" in arg for arg in binary_processes[-1].cmd)

    version_stdout, version_stderr, version_code = run_archivebox_cmd(
        ["version"],
        data_dir=initialized_archive,
        timeout=60,
        env=_runtime_env(initialized_archive, bootstrap_bin_dir),
    )
    assert version_code == 0, version_stderr
    assert name in version_stdout
    assert binary.version in version_stdout

    first_abspath.unlink()
    (initialized_archive / "lib" / "bin" / "rg").unlink(missing_ok=True)
    _link_real_binary(bootstrap_bin_dir, name, source="rg")

    rerun_stdout, rerun_stderr, rerun_code = run_archivebox_cmd(
        ["run", f"--binary-id={first_binary_id}"],
        data_dir=initialized_archive,
        timeout=120,
        env=_runtime_env(initialized_archive, bootstrap_bin_dir),
    )

    assert rerun_code == 0, rerun_stdout + rerun_stderr
    with use_archivebox_db(initialized_archive):
        recovered = Binary.objects.get(pk=first_binary_id)
        process_count = Process.objects.filter(process_type=Process.TypeChoices.BINARY).count()

    assert recovered.status == Binary.StatusChoices.INSTALLED
    assert recovered.version == binary.version
    assert Path(recovered.abspath).exists()
    assert Path(recovered.abspath).resolve() == Path(shutil.which("rg") or "").resolve()
    assert process_count >= 2


def test_missing_binary_request_stays_queued_then_recovers_when_provider_can_resolve(initialized_archive, tmp_path):
    name = f"abx-missing-rg-{uuid.uuid4().hex[:8]}"
    bootstrap_bin_dir = tmp_path / "realbin"
    provider_bin_dir = initialized_archive / "lib" / "env" / "bin"
    _link_real_binary(bootstrap_bin_dir, "uv")

    stdout, stderr, returncode = run_archivebox_cmd(
        ["run"],
        data_dir=initialized_archive,
        stdin=_binary_request(name),
        timeout=120,
        env=_runtime_env(initialized_archive, bootstrap_bin_dir),
    )

    assert returncode == 0, stderr
    assert any(record["type"] == "BinaryRequest" and record["name"] == name for record in parse_jsonl_output(stdout))

    with use_archivebox_db(initialized_archive):
        queued = Binary.objects.get(name=name)
        queued_id = str(queued.id)
        failed_process = Process.objects.filter(process_type=Process.TypeChoices.BINARY).latest("created_at")
        machine_config = Machine.objects.get(pk=queued.machine_id).config or {}

    assert queued.status == Binary.StatusChoices.QUEUED
    assert queued.abspath == ""
    assert queued.retry_at is not None
    assert failed_process.status == Process.StatusChoices.EXITED
    assert failed_process.exit_code == 1
    assert f"{name.upper().replace('-', '_')}_BINARY" not in machine_config
    assert not (provider_bin_dir / name).exists()

    _link_real_binary(provider_bin_dir, name, source="rg")

    recover_stdout, recover_stderr, recover_code = run_archivebox_cmd(
        ["run", f"--binary-id={queued_id}"],
        data_dir=initialized_archive,
        timeout=120,
        env=_runtime_env(initialized_archive, bootstrap_bin_dir),
    )

    assert recover_code == 0, recover_stdout + recover_stderr
    with use_archivebox_db(initialized_archive):
        recovered = Binary.objects.get(pk=queued_id)
        process_exit_codes = list(
            Process.objects.filter(process_type=Process.TypeChoices.BINARY).order_by("created_at").values_list("exit_code", flat=True),
        )

    assert recovered.status == Binary.StatusChoices.INSTALLED
    assert recovered.version
    assert Path(recovered.abspath).exists()
    assert Path(recovered.abspath) == provider_bin_dir / name
    assert process_exit_codes[-2:] == [1, 0]