feat: add fake acquisition pipeline
This commit is contained in:
@@ -0,0 +1,36 @@
|
|||||||
|
import time
|
||||||
|
|
||||||
|
import psycopg
|
||||||
|
|
||||||
|
# (state, currentStage) pairs the job moves through, in order, after claim.
|
||||||
|
_STAGES = [
|
||||||
|
("matched", "rank"),
|
||||||
|
("downloading", "download"),
|
||||||
|
("tagging", "tag"),
|
||||||
|
("imported", "import"),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> None:
|
||||||
|
with conn.cursor() as cur:
|
||||||
|
cur.execute(
|
||||||
|
'UPDATE "Job" SET state = %s, "currentStage" = %s, "updatedAt" = now() WHERE id = %s',
|
||||||
|
(state, stage, job_id),
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
|
||||||
|
def run_pipeline(conn: psycopg.Connection, job_id: str, stage_delay: float = 0.0) -> None:
|
||||||
|
"""Fake acquisition: walk the claimed job through to 'imported'."""
|
||||||
|
for state, stage in _STAGES:
|
||||||
|
if stage_delay:
|
||||||
|
time.sleep(stage_delay)
|
||||||
|
_set_state(conn, job_id, state, stage)
|
||||||
|
|
||||||
|
with conn.cursor() as cur:
|
||||||
|
cur.execute(
|
||||||
|
'UPDATE "Request" SET status = \'completed\' '
|
||||||
|
'WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)',
|
||||||
|
(job_id,),
|
||||||
|
)
|
||||||
|
conn.commit()
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
from lyra_worker.claim import claim_next
|
||||||
|
from lyra_worker.pipeline import run_pipeline
|
||||||
|
from tests.conftest import insert_request
|
||||||
|
|
||||||
|
|
||||||
|
def test_pipeline_drives_job_to_imported(conn):
|
||||||
|
job_id = insert_request(conn)
|
||||||
|
claim_next(conn)
|
||||||
|
|
||||||
|
run_pipeline(conn, job_id, stage_delay=0.0)
|
||||||
|
|
||||||
|
with conn.cursor() as cur:
|
||||||
|
cur.execute('SELECT state, "currentStage" FROM "Job" WHERE id = %s', (job_id,))
|
||||||
|
state, stage = cur.fetchone()
|
||||||
|
cur.execute('SELECT status FROM "Request" WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)', (job_id,))
|
||||||
|
req_status = cur.fetchone()[0]
|
||||||
|
|
||||||
|
assert state == "imported"
|
||||||
|
assert stage == "import"
|
||||||
|
assert req_status == "completed"
|
||||||
Reference in New Issue
Block a user