118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from game_data_capture.pipeline.jobs import JobLeases
|
|
from game_data_capture.storage.database import StorageDatabase
|
|
from game_data_capture.storage.repository import StorageError
|
|
|
|
|
|
@pytest.fixture
|
|
def database(tmp_path: Path) -> StorageDatabase:
|
|
database = StorageDatabase(
|
|
(tmp_path / "utility").resolve(), (Path(__file__).parents[3] / "migrations").resolve()
|
|
)
|
|
database.migrate()
|
|
return database
|
|
|
|
|
|
def _expire(database: StorageDatabase, job_id: str) -> None:
|
|
with database.connect() as connection:
|
|
connection.execute(
|
|
"UPDATE jobs SET lease_expires_at = ? WHERE job_id = ?",
|
|
((datetime.now(UTC) - timedelta(seconds=1)).isoformat(), job_id),
|
|
)
|
|
|
|
|
|
def test_job_validation_and_exclusive_ownership(database: StorageDatabase) -> None:
|
|
jobs = JobLeases(database)
|
|
for job_id, kind, attempts in (
|
|
("!", "ocr", 1),
|
|
("job", "bad", 1),
|
|
("job", "ocr", True),
|
|
("job", "ocr", 0),
|
|
):
|
|
with pytest.raises(ValueError):
|
|
jobs.create(job_id, kind, max_attempts=attempts)
|
|
jobs.create("job-1", "ocr")
|
|
for duration in (0, 301, True):
|
|
with pytest.raises(ValueError):
|
|
jobs.acquire("owner-1", duration)
|
|
with pytest.raises(ValueError):
|
|
jobs.acquire("!", 30)
|
|
assert jobs.acquire("owner-1", 30) == "job-1"
|
|
assert jobs.acquire("owner-2", 30) is None
|
|
with pytest.raises(StorageError):
|
|
jobs.renew("job-1", "owner-2", 30)
|
|
with pytest.raises(StorageError):
|
|
jobs.finish("job-1", "owner-2", "succeeded")
|
|
for duration in (0, 301, True):
|
|
with pytest.raises(ValueError):
|
|
jobs.renew("job-1", "owner-1", duration)
|
|
|
|
|
|
def test_expiry_recovery_exhaustion_and_terminal_jobs(database: StorageDatabase) -> None:
|
|
jobs = JobLeases(database)
|
|
jobs.create("retry", "ocr", max_attempts=2)
|
|
assert jobs.acquire("owner-1", 30) == "retry"
|
|
_expire(database, "retry")
|
|
with pytest.raises(StorageError):
|
|
jobs.renew("retry", "owner-1", 30)
|
|
with pytest.raises(StorageError):
|
|
jobs.finish("retry", "owner-1", "succeeded")
|
|
assert jobs.recover() == 1
|
|
with database.connect() as connection:
|
|
assert tuple(
|
|
connection.execute(
|
|
"SELECT state, lease_owner, lease_expires_at FROM jobs WHERE job_id='retry'"
|
|
).fetchone()
|
|
) == ("queued", None, None)
|
|
assert JobLeases(database).acquire("owner-2", 30) == "retry"
|
|
with database.connect() as connection:
|
|
assert tuple(
|
|
connection.execute(
|
|
"SELECT attempt, state, lease_owner FROM jobs WHERE job_id='retry'"
|
|
).fetchone()
|
|
) == (2, "running", "owner-2")
|
|
jobs.create("exhaust", "ocr", max_attempts=1)
|
|
assert jobs.acquire("owner-3", 30) == "exhaust"
|
|
_expire(database, "exhaust")
|
|
assert jobs.recover() == 1
|
|
with database.connect() as connection:
|
|
assert tuple(
|
|
connection.execute(
|
|
"SELECT state, attempt, lease_owner, lease_expires_at "
|
|
"FROM jobs WHERE job_id='exhaust'"
|
|
).fetchone()
|
|
) == ("failed", 1, None, None)
|
|
assert jobs.acquire("owner-4", 30) is None
|
|
for job_id, target in (("ok", "succeeded"), ("fail", "failed"), ("cancel", "cancelled")):
|
|
jobs.create(job_id, "ocr")
|
|
assert jobs.acquire("owner-5", 30) == job_id
|
|
jobs.finish(job_id, "owner-5", target)
|
|
with database.connect() as connection:
|
|
for job_id, target in (("ok", "succeeded"), ("fail", "failed"), ("cancel", "cancelled")):
|
|
assert tuple(
|
|
connection.execute(
|
|
"SELECT state, lease_owner, lease_expires_at FROM jobs WHERE job_id = ?",
|
|
(job_id,),
|
|
).fetchone()
|
|
) == (target, None, None)
|
|
|
|
|
|
def test_retry_wait_finish_rejected_without_changing_lease(database: StorageDatabase) -> None:
|
|
jobs = JobLeases(database)
|
|
jobs.create("job-1", "ocr")
|
|
jobs.acquire("owner-1", 30)
|
|
with pytest.raises(ValueError):
|
|
jobs.finish("job-1", "owner-1", "retry_wait")
|
|
with database.connect() as connection:
|
|
assert tuple(
|
|
connection.execute(
|
|
"SELECT state, lease_owner FROM jobs WHERE job_id='job-1'"
|
|
).fetchone()
|
|
) == ("running", "owner-1")
|