US02-03: Execute Jobs with Leases, Locks, and Recovery (#56)

tests/integration/test_worker.py

@@ -0,0 +1,156 @@
+"""Worker execution: success, failure, cooperative cancel, resume, fencing, and a
+randomized concurrency stress with no duplicate completion or lost events."""
+
+import threading
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from datetime import timedelta
+
+import pytest
+
+from photo_pipeline.db import create_db_engine, create_session_factory, run_migrations
+from photo_pipeline.jobs.worker import Worker
+from photo_pipeline.services.jobs import (
+    ItemState,
+    JobConflict,
+    JobService,
+    JobState,
+    _now,
+)
+
+
+@pytest.fixture
+def db_url(tmp_path):
+    url = f"sqlite:///{tmp_path / 'worker.db'}"
+    run_migrations(url)
+    return url
+
+
+@pytest.fixture
+def make_factory():
+    engines = []
+
+    def _make(url):
+        engine = create_db_engine(url)
+        engines.append(engine)
+        return create_session_factory(engine)
+
+    yield _make
+    for engine in engines:
+        engine.dispose()
+
+
+@pytest.fixture
+def sf(db_url, make_factory):
+    return make_factory(db_url)
+
+
+@pytest.fixture
+def jobs(sf):
+    return JobService(sf)
+
+
+def test_worker_processes_all_items(sf, jobs):
+    handled = []
+    worker = Worker(sf, {"scan": lambda item, ctx: handled.append(item)}, "w1")
+    job = jobs.enqueue("scan", items=["a", "b", "c"])
+    assert worker.run_once() == job["id"]
+    assert jobs.get(job["id"])["state"] == JobState.SUCCEEDED
+    assert sorted(handled) == ["a", "b", "c"]
+    assert jobs.progress(job["id"])["done"] == 3
+
+
+def test_handler_failure_fails_the_job(sf, jobs):
+    def handler(item, ctx):
+        if item == "bad":
+            raise ValueError("nope")
+
+    worker = Worker(sf, {"scan": handler}, "w1")
+    job = jobs.enqueue("scan", items=["ok", "bad"])
+    worker.run_once()
+    assert jobs.get(job["id"])["state"] == JobState.FAILED
+    assert jobs.progress(job["id"])["by_state"] == {
+        ItemState.SUCCEEDED: 1,
+        ItemState.FAILED: 1,
+    }
+
+
+def test_cooperative_cancellation_leaves_items_resumable(sf, jobs):
+    # Cancel arrives while the first item runs; the worker stops before the next.
+    def handler(item, ctx):
+        if item == "a":
+            ctx.service.cancel(ctx.job_id)
+
+    worker = Worker(sf, {"scan": handler}, "w1")
+    job = jobs.enqueue("scan", items=["a", "b"])
+    worker.run_once()
+    assert jobs.get(job["id"])["state"] == JobState.CANCELLED
+    by_state = jobs.progress(job["id"])["by_state"]
+    assert by_state.get(ItemState.SUCCEEDED) == 1  # "a" finished
+    assert by_state.get(ItemState.QUEUED) == 1  # "b" left resumable
+
+
+def test_fencing_rejects_superseded_worker(sf, jobs):
+    job = jobs.enqueue("scan", items=["a"])
+    stale = jobs.claim(["scan"], "old")
+    jobs.recover_stale(now=_now() + timedelta(hours=1))
+    fresh = jobs.claim(["scan"], "new")
+    assert fresh["fencing_token"] > stale["fencing_token"]
+    with pytest.raises(JobConflict):
+        jobs.set_item(job["id"], "a", ItemState.SUCCEEDED, fencing_token=stale["fencing_token"])
+    with pytest.raises(JobConflict):
+        jobs.transition(
+            job["id"], JobState.SUCCEEDED, fencing_token=stale["fencing_token"]
+        )
+
+
+def test_recovered_job_resumes_remaining_items(sf, jobs):
+    job = jobs.enqueue("scan", items=["a", "b"])
+    # Simulate a crashed attempt: "a" done, "b" left running, no finalize.
+    crashed = jobs.claim(["scan"], "dead")
+    token = crashed["fencing_token"]
+    jobs.set_item(job["id"], "a", ItemState.RUNNING, fencing_token=token)
+    jobs.set_item(job["id"], "a", ItemState.SUCCEEDED, fencing_token=token)
+    jobs.set_item(job["id"], "b", ItemState.RUNNING, fencing_token=token)
+    jobs.recover_stale(now=_now() + timedelta(hours=1))
+
+    handled = []
+    fresh_worker = Worker(sf, {"scan": lambda item, ctx: handled.append(item)}, "alive")
+    fresh_worker.run_once()
+    assert jobs.get(job["id"])["state"] == JobState.SUCCEEDED
+    assert handled == ["b"]  # "a" already succeeded, only "b" reprocessed
+    assert jobs.progress(job["id"])["by_state"][ItemState.SUCCEEDED] == 2
+
+
+def test_concurrent_workers_no_duplicate_completion_or_lost_events(db_url, make_factory):
+    service = JobService(make_factory(db_url))
+    job_ids = [service.enqueue("scan", items=["x", "y"])["id"] for _ in range(12)]
+
+    calls = Counter()
+    lock = threading.Lock()
+
+    def handler(item, ctx):
+        with lock:
+            calls[(ctx.job_id, item)] += 1
+
+    def drain(worker_index):
+        worker = Worker(make_factory(db_url), {"scan": handler}, f"w{worker_index}")
+        empties = 0
+        for _ in range(200):
+            if worker.run_once() is None:
+                empties += 1
+                if empties >= 2:
+                    break
+            else:
+                empties = 0
+
+    with ThreadPoolExecutor(max_workers=4) as pool:
+        list(pool.map(drain, range(4)))
+
+    # Every job completed exactly once; every item handled exactly once.
+    for job_id in job_ids:
+        assert service.get(job_id)["state"] == JobState.SUCCEEDED
+        types = [e["type"] for e in service.events(job_id)]
+        assert "claimed" in types and "state:succeeded" in types  # no lost events
+    assert all(count == 1 for count in calls.values())
+    assert len(calls) == len(job_ids) * 2