photoanalyzer/tests/integration/test_worker.py

"""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_a_handler_that_stops_itself_releases_the_lock(sf, jobs):
    """A handler may stop without anyone cancelling the *job* — an upload batch
    cancelled through its own API does exactly that. The job is still ``running``
    when it raises, so it has to reach ``cancelled`` through ``cancelling``; if that
    hop is skipped the transition is rejected and the lock is held forever."""
    from photo_pipeline.jobs.handlers import Cancelled

    def handler(item, ctx):
        raise Cancelled("the work this job wraps was stopped elsewhere")

    worker = Worker(sf, {"scan": handler}, "w1")
    job = jobs.enqueue("scan", lock="library_write", items=["a"])
    worker.run_once()

    assert jobs.get(job["id"])["state"] == JobState.CANCELLED
    assert jobs.progress(job["id"])["by_state"] == {ItemState.QUEUED: 1}  # resumable
    # The lane is free: the next job may be enqueued under the same lock.
    assert jobs.enqueue("scan", lock="library_write", items=["b"])["state"] == JobState.QUEUED


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