71 lines
3.7 KiB
Python
71 lines
3.7 KiB
Python
"""Opt-in real PostgreSQL queue races, using an explicitly disposable database."""
|
|
import asyncio
|
|
import os
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
from psycopg import AsyncConnection
|
|
from psycopg.conninfo import conninfo_to_dict
|
|
from psycopg.rows import dict_row
|
|
|
|
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
|
|
|
|
|
def test_concurrent_submission_claims_fencing_and_project_isolation():
|
|
dsn = os.getenv("CALCULATION_QUEUE_TEST_DSN")
|
|
if not dsn:
|
|
pytest.skip("Set CALCULATION_QUEUE_TEST_DSN to a disposable PostgreSQL database")
|
|
assert conninfo_to_dict(dsn)["dbname"].startswith("tjwater_calc_test_")
|
|
|
|
async def run():
|
|
project, other, owner = uuid4(), uuid4(), uuid4()
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
await conn.execute((Path(__file__).parents[2] / "resources/sql/026_calculation_jobs.sql").read_text())
|
|
|
|
async def submit(key, parameters=None, selected_project=project):
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
return await CalculationJobRepository(conn, selected_project).enqueue(
|
|
owner_user_id=owner, request_key=key, tool="run_burst_location",
|
|
parameters=parameters or {"x": 1}, binding={"native_session_id": "session"}, timeout_seconds=600)
|
|
|
|
# Four independent API connections, same logical submission.
|
|
submissions = await asyncio.gather(*(submit("same") for _ in range(4)))
|
|
assert len({str(row["job_id"]) for row in submissions}) == 1
|
|
with pytest.raises(HTTPException) as conflict:
|
|
await submit("same", {"x": 2})
|
|
assert conflict.value.status_code == 409
|
|
second = await submit("second")
|
|
foreign = await submit("foreign", selected_project=other)
|
|
|
|
async def claim(worker):
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
return await CalculationJobRepository(conn, project).claim(worker, 60)
|
|
|
|
claimed = await asyncio.gather(*(claim(f"worker-{index}") for index in range(4)))
|
|
rows = [row for row in claimed if row]
|
|
assert len(rows) == 2
|
|
assert len({str(row["job_id"]) for row in rows}) == 2
|
|
assert str(foreign["job_id"]) not in {str(row["job_id"]) for row in rows}
|
|
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
repo = CalculationJobRepository(conn, project)
|
|
assert await repo.get(foreign["job_id"], owner) is None
|
|
assert await repo.get(second["job_id"], uuid4()) is None
|
|
row = rows[0]
|
|
stale = {**row, "lease_generation": row["lease_generation"] - 1}
|
|
assert await repo.heartbeat(stale, row["lease_owner"], 60) is None
|
|
assert not await repo.finish(stale, row["lease_owner"], "succeeded", result={})
|
|
cancelled = await repo.cancel(row["job_id"], owner)
|
|
assert cancelled["status"] == "cancel_requested"
|
|
assert not await repo.finish(row, row["lease_owner"], "succeeded", result={})
|
|
assert await repo.finish(row, row["lease_owner"], "cancelled")
|
|
expired = rows[1]
|
|
await conn.execute("UPDATE workbench.calculation_jobs SET lease_expires_at=now()-interval '1 second' WHERE job_id=%s", (expired["job_id"],))
|
|
await repo.recover()
|
|
assert (await repo.get(expired["job_id"], owner))["status"] == "interrupted"
|
|
assert not await repo.finish(expired, expired["lease_owner"], "succeeded", result={})
|
|
assert await repo.claim("replacement", 60) is None
|
|
asyncio.run(run())
|