mirror of
https://github.com/ArchiveBox/ArchiveBox.git
synced 2026-09-14 11:06:13 +05:00
296 lines
9.8 KiB
Python
296 lines
9.8 KiB
Python
"""Tests for MAX/SIZE crawl limit config behavior."""
|
|
|
|
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from archivebox.core.models import Snapshot
|
|
from archivebox.crawls.models import Crawl
|
|
from archivebox.tests.conftest import cli_env, run_archivebox_cmd
|
|
from archivebox.tests.test_orm_helpers import use_archivebox_db
|
|
|
|
pytestmark = pytest.mark.django_db(transaction=True)
|
|
|
|
|
|
def test_create_snapshots_from_urls_respects_max_urls(admin_user):
|
|
crawl = Crawl.objects.create(
|
|
urls="\n".join(
|
|
[
|
|
"https://example.com/root",
|
|
"https://example.com/about",
|
|
"https://example.com/contact",
|
|
],
|
|
),
|
|
config={"CRAWL_MAX_URLS": 2},
|
|
created_by=admin_user,
|
|
)
|
|
|
|
created = crawl.create_snapshots_from_urls()
|
|
|
|
assert [snapshot.url for snapshot in created] == [
|
|
"https://example.com/root",
|
|
"https://example.com/about",
|
|
]
|
|
assert crawl.snapshot_set.count() == 2
|
|
assert crawl.remaining_snapshot_capacity() == 0
|
|
assert crawl.limit_stop_reason() == "crawl_max_urls"
|
|
assert crawl.add_url({"url": "https://example.com/extra", "depth": 1}) is False
|
|
|
|
|
|
def test_crawl_stop_reason_keeps_specific_limit_reason_over_lifecycle_fallback(admin_user):
|
|
crawl = Crawl.objects.create(
|
|
urls="\n".join(
|
|
[
|
|
"https://example.com/root",
|
|
"https://example.com/about",
|
|
],
|
|
),
|
|
config={"CRAWL_MAX_URLS": 1},
|
|
status=Crawl.StatusChoices.SEALED,
|
|
retry_at=None,
|
|
created_by=admin_user,
|
|
)
|
|
Snapshot.objects.create(
|
|
url="https://example.com/root",
|
|
crawl=crawl,
|
|
status=Snapshot.StatusChoices.SEALED,
|
|
timestamp="1700000000.011",
|
|
)
|
|
|
|
assert crawl.stop_reason() == "crawl_max_urls"
|
|
|
|
|
|
def test_enqueue_discovered_snapshots_refreshes_crawl_limits(tmp_path):
|
|
from archivebox.base_models.models import get_or_create_system_user_pk
|
|
from archivebox.crawls.models import Crawl
|
|
from archivebox.core.models import Snapshot
|
|
from archivebox.services.runner import CrawlRunner
|
|
|
|
crawl = Crawl.objects.create(
|
|
urls="https://example.com",
|
|
max_depth=0,
|
|
config={"CRAWL_MAX_URLS": 5},
|
|
created_by_id=get_or_create_system_user_pk(),
|
|
)
|
|
snapshot = Snapshot.objects.create(
|
|
url="https://example.com",
|
|
crawl=crawl,
|
|
status=Snapshot.StatusChoices.SEALED,
|
|
depth=0,
|
|
)
|
|
parser_dir = Path(snapshot.output_dir) / "parse_html_urls"
|
|
parser_dir.mkdir(parents=True, exist_ok=True)
|
|
(parser_dir / "urls.jsonl").write_text(
|
|
"\n".join(
|
|
[
|
|
json.dumps({"type": "Snapshot", "url": "https://example.com/child-a", "depth": 1}),
|
|
json.dumps({"type": "Snapshot", "url": "https://example.com/child-b", "depth": 1}),
|
|
"",
|
|
],
|
|
),
|
|
)
|
|
|
|
runner = CrawlRunner(crawl)
|
|
Crawl.objects.filter(id=crawl.id).update(max_depth=1)
|
|
payload = runner.load_snapshot_payload(str(snapshot.id))
|
|
|
|
asyncio.run(runner.enqueue_discovered_snapshots_from_outputs(payload))
|
|
|
|
child_snapshots = list(crawl.snapshot_set.filter(depth=1).order_by("url").values_list("url", "status"))
|
|
assert child_snapshots == [
|
|
("https://example.com/child-a", Snapshot.StatusChoices.QUEUED),
|
|
("https://example.com/child-b", Snapshot.StatusChoices.QUEUED),
|
|
]
|
|
|
|
|
|
def test_run_snapshot_seals_descendant_when_crawl_max_size_is_reached(tmp_path):
|
|
from abx_dl.events import CrawlStartEvent, SnapshotEvent
|
|
from archivebox.base_models.models import get_or_create_system_user_pk
|
|
from archivebox.crawls.models import Crawl
|
|
from archivebox.core.models import Snapshot
|
|
from archivebox.services.runner import CrawlRunner
|
|
|
|
crawl = Crawl.objects.create(
|
|
urls="https://example.com",
|
|
config={
|
|
"LIB_DIR": str(tmp_path / "lib"),
|
|
"PLUGINS": "__archivebox_test_no_plugins__",
|
|
"CHROME_BINARY": "",
|
|
"CRAWL_MAX_SIZE": 16,
|
|
},
|
|
created_by_id=get_or_create_system_user_pk(),
|
|
)
|
|
root = Snapshot.objects.create(
|
|
url="https://example.com",
|
|
crawl=crawl,
|
|
depth=0,
|
|
status=Snapshot.StatusChoices.SEALED,
|
|
)
|
|
child = Snapshot.objects.create(
|
|
url="https://example.com/child",
|
|
crawl=crawl,
|
|
depth=1,
|
|
parent_snapshot=root,
|
|
status=Snapshot.StatusChoices.QUEUED,
|
|
)
|
|
state_dir = Path(crawl.output_dir) / ".abx-dl"
|
|
state_dir.mkdir(parents=True, exist_ok=True)
|
|
(state_dir / "limits.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"admitted_snapshot_ids": [str(root.id), str(child.id)],
|
|
"counted_event_ids": ["proc-1"],
|
|
"total_size": 32,
|
|
"stop_reason": "crawl_max_size",
|
|
},
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
runner = CrawlRunner(crawl)
|
|
runner.load_run_state()
|
|
|
|
async def run_child_from_crawl_start() -> list[SnapshotEvent]:
|
|
async def on_crawl_start(_event: CrawlStartEvent) -> None:
|
|
await runner.run_snapshot(str(child.id))
|
|
|
|
runner.bus.on(CrawlStartEvent, on_crawl_start)
|
|
crawl_start = runner.bus.emit(
|
|
CrawlStartEvent(
|
|
url=root.url,
|
|
snapshot_id=str(root.id),
|
|
output_dir=str(crawl.output_dir),
|
|
event_timeout=30,
|
|
event_handler_timeout=30,
|
|
),
|
|
)
|
|
await crawl_start.now()
|
|
await crawl_start.event_results_list()
|
|
await runner.bus.wait_until_idle()
|
|
return await runner.bus.filter(SnapshotEvent, child_of=crawl_start, past=True)
|
|
|
|
snapshot_events = asyncio.run(run_child_from_crawl_start())
|
|
|
|
child.refresh_from_db()
|
|
assert child.status == Snapshot.StatusChoices.SEALED
|
|
assert child.retry_at is None
|
|
assert snapshot_events == []
|
|
|
|
|
|
def test_seal_snapshot_cancels_queued_descendants_after_crawl_max_size():
|
|
from archivebox.base_models.models import get_or_create_system_user_pk
|
|
from archivebox.crawls.models import Crawl
|
|
from archivebox.core.models import Snapshot
|
|
from archivebox.services.snapshot_service import SnapshotService
|
|
from abx_dl.events import SnapshotCompletedEvent
|
|
from abx_dl.orchestrator import create_bus
|
|
|
|
crawl = Crawl.objects.create(
|
|
urls="https://example.com",
|
|
created_by_id=get_or_create_system_user_pk(),
|
|
config={"CRAWL_MAX_SIZE": 16},
|
|
)
|
|
root = Snapshot.objects.create(
|
|
url="https://example.com",
|
|
crawl=crawl,
|
|
status=Snapshot.StatusChoices.STARTED,
|
|
)
|
|
child = Snapshot.objects.create(
|
|
url="https://example.com/child",
|
|
crawl=crawl,
|
|
depth=1,
|
|
parent_snapshot_id=root.id,
|
|
status=Snapshot.StatusChoices.QUEUED,
|
|
)
|
|
|
|
state_dir = Path(crawl.output_dir) / ".abx-dl"
|
|
state_dir.mkdir(parents=True, exist_ok=True)
|
|
(state_dir / "limits.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"admitted_snapshot_ids": [str(root.id), str(child.id)],
|
|
"counted_event_ids": ["proc-1"],
|
|
"total_size": 32,
|
|
"stop_reason": "crawl_max_size",
|
|
},
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
bus = create_bus(name=f"test_snapshot_limit_cancel_{str(crawl.id).replace('-', '_')}")
|
|
service = SnapshotService(bus, crawl_id=str(crawl.id), schedule_snapshot=lambda snapshot_id: None)
|
|
try:
|
|
|
|
async def emit_event() -> None:
|
|
await service.on_SnapshotCompletedEvent(
|
|
SnapshotCompletedEvent(
|
|
url=root.url,
|
|
snapshot_id=str(root.id),
|
|
output_dir=str(root.output_dir),
|
|
),
|
|
)
|
|
|
|
asyncio.run(emit_event())
|
|
finally:
|
|
asyncio.run(bus.wait_until_idle())
|
|
asyncio.run(bus.destroy())
|
|
|
|
root.refresh_from_db()
|
|
child.refresh_from_db()
|
|
assert root.status == Snapshot.StatusChoices.SEALED
|
|
assert child.status == Snapshot.StatusChoices.SEALED
|
|
assert child.retry_at is None
|
|
|
|
|
|
def test_recursive_crawl_respects_max_urls(tmp_path, initialized_archive, recursive_test_site):
|
|
"""Test that recursive discovery stops creating snapshots at max_urls."""
|
|
env = cli_env(disable_extractors=True)
|
|
|
|
env = env.copy()
|
|
env.update(
|
|
{
|
|
"URL_ALLOWLIST": r"127\.0\.0\.1[:/].*",
|
|
"SAVE_WGET": "true",
|
|
"USE_CHROME": "false",
|
|
"USE_COLOR": "false",
|
|
"SHOW_PROGRESS": "false",
|
|
},
|
|
)
|
|
|
|
result = run_archivebox_cmd(
|
|
[
|
|
"add",
|
|
"--depth=2",
|
|
"--max-urls=4",
|
|
"--plugins=wget,parse_html_urls",
|
|
recursive_test_site["root_url"],
|
|
],
|
|
env=env,
|
|
timeout=120,
|
|
)
|
|
stdout, stderr = result.stdout, result.stderr
|
|
|
|
if stderr:
|
|
print(f"\n=== STDERR ===\n{stderr}\n=== END STDERR ===\n")
|
|
if stdout:
|
|
print(f"\n=== STDOUT (last 2000 chars) ===\n{stdout[-2000:]}\n=== END STDOUT ===\n")
|
|
|
|
assert result.returncode == 0, result.stderr
|
|
|
|
with use_archivebox_db(tmp_path):
|
|
crawl_obj = Crawl.objects.order_by("-created_at").first()
|
|
crawl = (crawl_obj.max_depth, crawl_obj.config["CRAWL_MAX_URLS"]) if crawl_obj else None
|
|
snapshot_rows = list(Snapshot.objects.order_by("depth", "url").values_list("url", "depth", "parent_snapshot_id"))
|
|
depth_counts = {
|
|
depth: Snapshot.objects.filter(depth=depth).count() for depth in set(Snapshot.objects.values_list("depth", flat=True))
|
|
}
|
|
|
|
assert crawl == (2, 4)
|
|
assert len(snapshot_rows) == 4
|
|
assert depth_counts.get(0, 0) == 1
|
|
assert depth_counts.get(1, 0) == 3
|
|
assert depth_counts.get(2, 0) == 0
|
|
assert set(recursive_test_site["child_urls"]).issubset({url for url, depth, _parent in snapshot_rows if depth == 1})
|