Blame for tests/test_queue.py

"""Concurrency and lifecycle guarantees for the inference queue."""

from __future__ import annotations

import asyncio
import threading
import time
from contextvars import ContextVar

import pytest

from core.errors import GatewayError, QueueCapacityError
from core.queue import InferenceQueue


async def test_queue_executes_only_one_operation_at_a_time() -> None:
    queue = InferenceQueue(capacity=4, timeout_seconds=5)
    active = 0
    maximum_active = 0
    lock = threading.Lock()

    def operation(value: int) -> int:
        nonlocal active, maximum_active
        with lock:
            active += 1
            maximum_active = max(maximum_active, active)
        time.sleep(0.02)
        with lock:
            active -= 1
        return value

    results = await asyncio.gather(
        *(queue.submit(str(index), "test", lambda index=index: operation(index)) for index in range(3))
    )
    await queue.stop()

    assert results == [0, 1, 2]
    assert maximum_active == 1


async def test_queue_timeout_cancels_job_before_inference() -> None:
    queue = InferenceQueue(capacity=2, timeout_seconds=0.01)
    queued_job_ran = False

    async def submit_first() -> None:
        await queue.submit("first", "slow", lambda: time.sleep(0.04))

    def queued_operation() -> None:
        nonlocal queued_job_ran
        queued_job_ran = True

    first = asyncio.create_task(submit_first())
    await asyncio.sleep(0.005)

    with pytest.raises(GatewayError, match="timed out") as captured:
        await queue.submit("second", "queued", queued_operation)

    assert captured.value.status_code == 504
    assert captured.value.code == "queue_timeout"
    record = queue.job("second")
    assert record is not None
    assert record["status"] == "timeout"
    await first
    await asyncio.sleep(0)
    assert queued_job_ran is False
    assert queue.active is False
    await queue.stop()


async def test_cancellation_waits_for_started_operation() -> None:
    queue = InferenceQueue(capacity=2, timeout_seconds=1)
    started = threading.Event()
    finished = threading.Event()
    output = []

    def operation() -> int:
        started.set()
        time.sleep(0.04)
        output.append("complete")
        finished.set()
        return 7

    task = asyncio.create_task(queue.submit("request", "slow", operation))
    assert await asyncio.to_thread(started.wait, 0.1)

    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task

    assert finished.is_set()
    assert output == ["complete"]
    assert queue.job("request")["status"] == "succeeded"

    result = await queue.submit("next", "test", lambda: 42)
    assert result == 42
    await queue.stop()


async def test_cancellation_before_start_marks_job_cancelled_and_does_not_run() -> None:
    queue = InferenceQueue(capacity=2, timeout_seconds=1)
    blocker_started = threading.Event()
    blocker_finished = threading.Event()
    queued_ran = False

    def blocker() -> None:
        blocker_started.set()
        time.sleep(0.05)
        blocker_finished.set()

    def queued_operation() -> None:
        nonlocal queued_ran
        queued_ran = True

    first = asyncio.create_task(queue.submit("blocker", "slow", blocker))
    assert await asyncio.to_thread(blocker_started.wait, 0.1)

    second = asyncio.create_task(queue.submit("cancelled", "queued", queued_operation))
    await asyncio.sleep(0.005)
    second.cancel()

    with pytest.raises(asyncio.CancelledError):
        await second
    await first

    assert blocker_finished.is_set()
    assert queued_ran is False
    record = queue.job("cancelled")
    assert record is not None
    assert record["status"] == "cancelled"
    assert record["finished_at"] is not None
    await queue.stop()


async def test_queue_preserves_request_context_for_worker() -> None:
    queue = InferenceQueue(capacity=1, timeout_seconds=5)
    request_context: ContextVar[str] = ContextVar("test_request", default="missing")
    token = request_context.set("request-identity")
    try:
        result = await queue.submit("request", "context", request_context.get)
    finally:
        request_context.reset(token)
        await queue.stop()

    assert result == "request-identity"


async def test_duplicate_request_id_is_rejected_without_overwriting_history() -> None:
    queue = InferenceQueue(capacity=2, timeout_seconds=1, history_size=10)

    assert await queue.submit("duplicate", "first", lambda: 1) == 1

    with pytest.raises(GatewayError) as captured:
        await queue.submit("duplicate", "second", lambda: 2)

    assert captured.value.status_code == 409
    assert captured.value.code == "request_id_conflict"
    record = queue.job("duplicate")
    assert record is not None
    assert record["label"] == "first"
    assert record["status"] == "succeeded"
    await queue.stop()


async def test_queue_history_is_bounded_and_newest_first() -> None:
    queue = InferenceQueue(capacity=2, timeout_seconds=1, history_size=2)

    await queue.submit("one", "first", lambda: 1)
    await queue.submit("two", "second", lambda: 2)
    await queue.submit("three", "third", lambda: 3)

    records = queue.jobs()
    assert [record["request_id"] for record in records] == ["three", "two"]
    assert queue.job("one") is None
    await queue.stop()


async def test_operation_failure_is_recorded() -> None:
    queue = InferenceQueue(capacity=2, timeout_seconds=1)

    def fail() -> None:
        raise RuntimeError("model failed")

    with pytest.raises(RuntimeError, match="model failed"):
        await queue.submit("failed", "failure", fail)

    record = queue.job("failed")
    assert record is not None
    assert record["status"] == "failed"
    assert record["error_code"] == "inference_failed"
    assert record["error_message"] == "model failed"
    assert record["finished_at"] is not None
    await queue.stop()


async def test_queue_full_rejects_without_running_the_rejected_operation() -> None:
    queue = InferenceQueue(capacity=1, timeout_seconds=1)
    started = threading.Event()
    release = threading.Event()
    rejected_ran = False

    def blocker() -> None:
        started.set()
        release.wait(0.2)

    def rejected_operation() -> None:
        nonlocal rejected_ran
        rejected_ran = True

    first = asyncio.create_task(queue.submit("running", "blocker", blocker))
    assert await asyncio.to_thread(started.wait, 0.1)

    queued = asyncio.create_task(queue.submit("queued", "waiting", lambda: 2))
    await asyncio.sleep(0.005)

    with pytest.raises(QueueCapacityError):
        await queue.submit("rejected", "overflow", rejected_operation)

    rejected_record = queue.job("rejected")
    assert rejected_record is not None
    assert rejected_record["status"] == "rejected"
    assert rejected_record["error_code"] == "queue_full"

    release.set()
    assert await first is None
    assert await queued == 2
    assert rejected_ran is False
    await queue.stop()


async def test_shutdown_cancels_queued_jobs_but_waits_for_active_operation() -> None:
    queue = InferenceQueue(capacity=2, timeout_seconds=1)
    started = threading.Event()
    release = threading.Event()

    def active_operation() -> int:
        started.set()
        release.wait(0.2)
        return 1

    active_task = asyncio.create_task(queue.submit("active", "active", active_operation))
    assert await asyncio.to_thread(started.wait, 0.1)

    queued_task = asyncio.create_task(queue.submit("queued", "queued", lambda: 2))
    await asyncio.sleep(0.005)

    stop_task = asyncio.create_task(queue.stop(timeout_seconds=0.5))
    await asyncio.sleep(0.005)

    queued_record = queue.job("queued")
    assert queued_record is not None
    assert queued_record["status"] == "cancelled"

    release.set()
    assert await active_task == 1
    assert await stop_task is True

    with pytest.raises(asyncio.CancelledError):
        await queued_task


async def test_shutdown_timeout_resolves_active_request_and_records_timeout() -> None:
    queue = InferenceQueue(capacity=1, timeout_seconds=1)
    started = threading.Event()
    release = threading.Event()

    def active_operation() -> int:
        started.set()
        release.wait(0.5)
        return 9

    submit_task = asyncio.create_task(queue.submit("shutdown-active", "slow", active_operation))
    assert await asyncio.to_thread(started.wait, 0.1)

    assert await queue.stop(timeout_seconds=0.01) is False

    with pytest.raises(GatewayError) as captured:
        await asyncio.wait_for(submit_task, timeout=0.2)

    assert captured.value.status_code == 503
    assert captured.value.code == "shutdown_timeout"
    record = queue.job("shutdown-active")
    assert record is not None
    assert record["status"] == "timeout"
    assert record["error_code"] == "shutdown_timeout"
    assert record["finished_at"] is not None

    # The Python/native worker thread is not force-killed by asyncio cancellation.
    release.set()



async def test_cancellation_remains_primary_when_started_operation_fails() -> None:
    queue = InferenceQueue(capacity=1, timeout_seconds=1)
    started = threading.Event()
    release = threading.Event()

    def fail_after_release() -> int:
        started.set()
        release.wait(0.5)
        raise RuntimeError("model failed after cancellation")

    task = asyncio.create_task(
        queue.submit("cancel-then-fail", "failure", fail_after_release)
    )
    assert await asyncio.to_thread(started.wait, 0.1)
    # Let submit resume from the start signal and enter its result wait.
    await asyncio.sleep(0.01)

    task.cancel()
    release.set()
    with pytest.raises(asyncio.CancelledError):
        await task

    record = queue.job("cancel-then-fail")
    assert record is not None
    assert record["status"] == "failed"
    assert record["error_code"] == "inference_failed"
    assert record["error_message"] == "model failed after cancellation"
    assert record["finished_at"] is not None
    await queue.stop()



async def test_shutdown_timeout_remains_failed_on_repeated_stop() -> None:
    queue = InferenceQueue(capacity=1, timeout_seconds=1)
    started = threading.Event()
    release = threading.Event()

    def active_operation() -> int:
        started.set()
        release.wait(0.5)
        return 1

    submit_task = asyncio.create_task(
        queue.submit("repeat-stop", "slow", active_operation)
    )
    assert await asyncio.to_thread(started.wait, 0.1)

    assert await queue.stop(timeout_seconds=0.01) is False
    assert await queue.stop(timeout_seconds=0.01) is False

    with pytest.raises(GatewayError) as captured:
        await asyncio.wait_for(submit_task, timeout=0.2)
    assert captured.value.code == "shutdown_timeout"

    record = queue.job("repeat-stop")
    assert record is not None
    assert record["status"] == "timeout"
    release.set()



async def test_cancelled_caller_during_shutdown_timeout_keeps_consistent_state() -> None:
    queue = InferenceQueue(capacity=1, timeout_seconds=1)
    started = threading.Event()
    release = threading.Event()

    def active_operation() -> int:
        started.set()
        release.wait(0.5)
        return 5

    task = asyncio.create_task(
        queue.submit("cancel-during-stop", "slow", active_operation)
    )
    assert await asyncio.to_thread(started.wait, 0.1)
    await asyncio.sleep(0.01)

    task.cancel()
    stop_task = asyncio.create_task(queue.stop(timeout_seconds=0.01))
    assert await stop_task is False

    with pytest.raises(asyncio.CancelledError):
        await asyncio.wait_for(task, timeout=0.2)

    record = queue.job("cancel-during-stop")
    assert record is not None
    assert record["status"] == "timeout"
    assert record["error_code"] == "shutdown_timeout"
    assert record["finished_at"] is not None
    release.set()


async def test_submit_racing_shutdown_cannot_enqueue_behind_worker_sentinel() -> None:
    queue = InferenceQueue(capacity=1, timeout_seconds=1)
    entered = asyncio.Event()

    async def submit_job() -> None:
        entered.set()
        await queue.submit("submit-stop-race", "race", lambda: 1)

    # Hold admission while both tasks queue for the lifecycle lock. Whichever
    # acquires it first, shutdown must either cancel the job or reject admission.
    await queue._lifecycle_lock.acquire()
    submit_task = asyncio.create_task(submit_job())
    await entered.wait()
    stop_task = asyncio.create_task(queue.stop(timeout_seconds=0.2))
    await asyncio.sleep(0)
    queue._lifecycle_lock.release()

    stop_result = await asyncio.wait_for(stop_task, timeout=0.5)
    assert stop_result is True
    try:
        await asyncio.wait_for(submit_task, timeout=0.5)
    except (asyncio.CancelledError, GatewayError):
        pass

    assert queue.stopping is True
    assert queue.depth == 0


async def test_concurrent_stop_calls_enqueue_only_one_sentinel() -> None:
    queue = InferenceQueue(capacity=2, timeout_seconds=1)
    started = threading.Event()
    release = threading.Event()

    def operation() -> int:
        started.set()
        release.wait(0.5)
        return 3

    submit_task = asyncio.create_task(queue.submit("concurrent-stop", "slow", operation))
    assert await asyncio.to_thread(started.wait, 0.1)

    first_stop = asyncio.create_task(queue.stop(timeout_seconds=0.5))
    second_stop = asyncio.create_task(queue.stop(timeout_seconds=0.5))
    await asyncio.sleep(0.01)
    release.set()

    assert await asyncio.wait_for(submit_task, timeout=0.5) == 3
    results = await asyncio.gather(first_stop, second_stop)
    assert results == [True, True]
    assert queue.depth == 0
    assert queue.stopping is True