167 lines
5.0 KiB
Python
167 lines
5.0 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime, timedelta
|
|
|
|
from app.db import Database
|
|
|
|
|
|
def capture_config() -> dict[str, int | bool]:
|
|
return {
|
|
"include_external_assets": True,
|
|
"max_pages": 20,
|
|
"max_depth": 3,
|
|
"max_bytes": 1_000_000,
|
|
"max_duration_seconds": 60,
|
|
"retention_days": 1,
|
|
}
|
|
|
|
|
|
def test_sqlite_lease_claims_one_queued_job(tmp_path) -> None:
|
|
database = Database(tmp_path / "siteharbor.sqlite3")
|
|
database.initialize()
|
|
job = database.create_job(
|
|
"https://example.com",
|
|
"https://example.com/",
|
|
capture_config(),
|
|
)
|
|
|
|
claimed = database.claim_next("worker-a", lease_seconds=30)
|
|
|
|
assert claimed is not None
|
|
assert claimed["id"] == job["id"]
|
|
assert claimed["state"] == "running"
|
|
assert database.claim_next("worker-b", lease_seconds=30) is None
|
|
|
|
assert database.update_progress(
|
|
job["id"],
|
|
"worker-a",
|
|
30,
|
|
phase="Packaging",
|
|
message="Creating archive",
|
|
stats={
|
|
"pages_found": 1,
|
|
"pages_fetched": 1,
|
|
"assets_found": 0,
|
|
"assets_fetched": 0,
|
|
"bytes_downloaded": 10,
|
|
"files_written": 1,
|
|
},
|
|
warnings=[],
|
|
)
|
|
assert database.mark_ready(
|
|
job["id"],
|
|
"worker-a",
|
|
artifact_path=str(tmp_path / "capture.zip"),
|
|
report_path=str(tmp_path / "capture.json"),
|
|
archive_size=10,
|
|
entry_point="example.com/index.html",
|
|
stats={
|
|
"pages_found": 1,
|
|
"pages_fetched": 1,
|
|
"assets_found": 0,
|
|
"assets_fetched": 0,
|
|
"bytes_downloaded": 10,
|
|
"files_written": 1,
|
|
},
|
|
warnings=[],
|
|
retention_days=30,
|
|
)
|
|
ready = database.get_job(job["id"])
|
|
assert ready is not None
|
|
assert ready["state"] == "ready"
|
|
assert ready["expires_at"] > ready["completed_at"]
|
|
expires_at = datetime.fromisoformat(ready["expires_at"])
|
|
completed_at = datetime.fromisoformat(ready["completed_at"])
|
|
assert expires_at - completed_at <= timedelta(days=2)
|
|
|
|
|
|
def test_worker_cannot_cancel_without_a_user_cancellation_request(tmp_path) -> None:
|
|
database = Database(tmp_path / "siteharbor.sqlite3")
|
|
database.initialize()
|
|
job = database.create_job(
|
|
"https://example.com",
|
|
"https://example.com/",
|
|
{
|
|
**capture_config(),
|
|
"include_external_assets": False,
|
|
"max_pages": 1,
|
|
"max_depth": 0,
|
|
"max_bytes": 1_000,
|
|
"max_duration_seconds": 30,
|
|
},
|
|
)
|
|
assert database.claim_next("worker-a", lease_seconds=30)
|
|
|
|
assert not database.mark_cancelled(job["id"], "worker-a", "Lease was lost")
|
|
running = database.get_job(job["id"])
|
|
assert running is not None
|
|
assert running["state"] == "running"
|
|
|
|
database.request_cancel(job["id"])
|
|
assert database.mark_cancelled(job["id"], "worker-a", "Cancelled by user")
|
|
cancelled = database.get_job(job["id"])
|
|
assert cancelled is not None
|
|
assert cancelled["state"] == "cancelled"
|
|
|
|
|
|
def test_expired_cancel_requested_job_is_terminalized_before_reclaim(tmp_path) -> None:
|
|
database = Database(tmp_path / "siteharbor.sqlite3")
|
|
database.initialize()
|
|
job = database.create_job(
|
|
"https://example.com",
|
|
"https://example.com/",
|
|
{
|
|
**capture_config(),
|
|
"include_external_assets": False,
|
|
"max_pages": 1,
|
|
"max_depth": 0,
|
|
"max_bytes": 1_000,
|
|
"max_duration_seconds": 30,
|
|
},
|
|
)
|
|
assert database.claim_next("worker-a", lease_seconds=30)
|
|
with database.connection() as connection:
|
|
connection.execute(
|
|
"""
|
|
UPDATE jobs
|
|
SET cancel_requested = 1, lease_expires_at = '2000-01-01T00:00:00+00:00'
|
|
WHERE id = ?
|
|
""",
|
|
(job["id"],),
|
|
)
|
|
|
|
assert database.claim_next("worker-b", lease_seconds=30) is None
|
|
cancelled = database.get_job(job["id"])
|
|
assert cancelled is not None
|
|
assert cancelled["state"] == "cancelled"
|
|
|
|
|
|
def test_owner_token_scopes_jobs_and_public_stats(tmp_path) -> None:
|
|
database = Database(tmp_path / "siteharbor.sqlite3")
|
|
database.initialize()
|
|
first_owner = "a" * 32
|
|
second_owner = "b" * 32
|
|
first = database.create_job(
|
|
"https://one.example",
|
|
"https://one.example/",
|
|
capture_config(),
|
|
first_owner,
|
|
)
|
|
second = database.create_job(
|
|
"https://two.example",
|
|
"https://two.example/",
|
|
capture_config(),
|
|
second_owner,
|
|
)
|
|
|
|
assert database.get_owned_job(first["id"], first_owner) is not None
|
|
assert database.get_owned_job(first["id"], second_owner) is None
|
|
assert [job["id"] for job in database.list_owned_jobs(first_owner)] == [first["id"]]
|
|
assert [job["id"] for job in database.list_owned_jobs(second_owner)] == [second["id"]]
|
|
assert database.public_stats() == {
|
|
"websites_cloned": 0,
|
|
"bytes_scraped": 0,
|
|
"active_crawls": 2,
|
|
"files_archived": 0,
|
|
}
|