STT-Assistant/tests/test_inference.py

"""Tests for the inference queue (Phase B)."""

import threading
import time

from stt_daemon.assistant.inference import InferenceJob, InferenceQueue


def _wait_in_flight(runner: "_JobRunner", timeout: float = 1.0) -> None:
    """Block until a job is executing (or fail the caller's timeout)."""
    deadline = time.monotonic() + timeout
    while not runner.in_flight():
        assert time.monotonic() < deadline, "no job started in time"
        time.sleep(0.001)


def _wait_idle(queue: InferenceQueue, timeout: float = 2.0) -> None:
    """Block until nothing is queued and no job is in-flight."""
    deadline = time.monotonic() + timeout
    while queue.depth() > 0 or queue.in_flight():
        assert time.monotonic() < deadline, "queue did not drain in time"
        time.sleep(0.001)


class _JobRunner:
    """Helper to track job execution timing and detect concurrency."""

    def __init__(self):
        self.events: list[tuple[str, str]] = []  # [(event, job_name), ...]
        self.lock = threading.Lock()

    def job(self, name: str, block_until: threading.Event | None = None):
        """Return a callable that records start/end and optionally blocks."""

        def run():
            with self.lock:
                self.events.append(("start", name))
            if block_until is not None:
                block_until.wait()
            with self.lock:
                self.events.append(("end", name))

        return run

    def in_flight(self) -> str | None:
        """Return the name of the job currently running, or None."""
        with self.lock:
            for i in range(len(self.events) - 1, -1, -1):
                event, name = self.events[i]
                if event == "start":
                    return name
                if event == "end":
                    return None
        return None


def test_jobs_run_one_at_a_time():
    """Max concurrency is 1: job 2 does not start until job 1 ends."""
    queue = InferenceQueue(size=4)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()

    job1 = InferenceJob(name="job1", run=runner.job("job1", unblock))
    job2 = InferenceJob(name="job2", run=runner.job("job2"))

    queue.submit(job1)
    queue.submit(job2)

    # Give job1 time to start
    import time
    time.sleep(0.01)

    # job2 must not have started yet
    assert runner.in_flight() == "job1"

    # Unblock job1
    unblock.set()
    time.sleep(0.05)

    # Now job2 should be running or done
    assert any(event[1] == "job2" for event in runner.events)

    queue.stop()


def test_priority_jobs_inserted_at_front():
    """Priority job (priority < 0) submitted while normal job queued runs next."""
    queue = InferenceQueue(size=4)
    runner = _JobRunner()
    unblock1 = threading.Event()
    unblock_cmd = threading.Event()

    queue.start()

    job1 = InferenceJob(name="normal", run=runner.job("normal", unblock1), priority=0)
    job_cmd = InferenceJob(name="command", run=runner.job("command", unblock_cmd), priority=-1)
    job3 = InferenceJob(name="chat2", run=runner.job("chat2"), priority=0)

    queue.submit(job1)
    queue.submit(job3)
    queue.submit(job_cmd)

    import time
    time.sleep(0.1)

    # Release all blocks to complete all jobs
    unblock1.set()
    unblock_cmd.set()
    time.sleep(0.1)

    # Check that command ran before chat2 (priority ordering)
    names = [name for _, name in runner.events if _ == "end"]
    assert "command" in names
    assert "chat2" in names
    assert "normal" in names
    # Priority guarantees: command should run before chat2
    cmd_index = names.index("command")
    chat_index = names.index("chat2")
    assert cmd_index < chat_index

    queue.stop()


def test_priority_job_rejected_when_queue_full():
    """Overflow: priority job is rejected (submit returns False)."""
    queue = InferenceQueue(size=2)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()

    job1 = InferenceJob(name="job1", run=runner.job("job1", unblock), priority=0)
    job2 = InferenceJob(name="job2", run=runner.job("job2"), priority=0)
    job_cmd = InferenceJob(name="command", run=runner.job("command"), priority=-1)

    assert queue.submit(job1) is True
    assert queue.submit(job2) is True
    # Queue is now full (size=2: job1 in-flight, job2 queued)
    assert queue.submit(job_cmd) is False

    queue.stop()


def test_normal_job_evicts_oldest_on_overflow():
    """Overflow: normal job drops the oldest normal job."""
    queue = InferenceQueue(size=2)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()

    job1 = InferenceJob(name="job1", run=runner.job("job1", unblock), priority=0)
    job2 = InferenceJob(name="job2", run=runner.job("job2"), priority=0)
    job3 = InferenceJob(name="job3", run=runner.job("job3"), priority=0)

    assert queue.submit(job1) is True
    import time
    time.sleep(0.05)

    # Ensure job1 is in-flight
    while not runner.in_flight():
        time.sleep(0.001)

    # Now queue is: job1 in-flight, job2 queued, room for 1 more
    assert queue.submit(job2) is True
    # Queue is now full (in-flight + 1 queued)
    # job3 submits, evicts job2 (oldest queued normal job)
    assert queue.submit(job3) is True

    unblock.set()
    time.sleep(0.1)

    # Only job1 and job3 should have run; job2 was evicted
    names = [name for _, name in runner.events if _ == "end"]
    assert "job2" not in names
    assert set(names) == {"job1", "job3"}

    queue.stop()


def test_drop_removes_matching_queued_jobs():
    """drop(predicate) removes queued jobs matching predicate, returns count."""
    queue = InferenceQueue(size=8)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()

    job1 = InferenceJob(name="chat", run=runner.job("chat1", unblock), priority=0)
    job2 = InferenceJob(name="chat", run=runner.job("chat2"), priority=0)
    job3 = InferenceJob(name="cmd", run=runner.job("cmd"), priority=-1)

    queue.submit(job1)
    queue.submit(job2)
    queue.submit(job3)

    import time
    time.sleep(0.05)

    # Give time for job1 to start in-flight
    while not runner.in_flight():
        time.sleep(0.001)

    # Drop all chat jobs; job1 is in-flight, job2 is queued
    count = queue.drop(lambda j: j.name == "chat")
    assert count == 1  # job2 is queued; job1 is in-flight and not dropped

    unblock.set()
    time.sleep(0.1)

    # Only job1 (in-flight, not dropped) and job3 should run
    names = [name for _, name in runner.events if _ == "end"]
    assert set(names) == {"chat1", "cmd"}

    queue.stop()


def test_stop_discards_queued_and_joins():
    """stop() discards queued jobs, lets in-flight finish, then joins."""
    queue = InferenceQueue(size=4)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()

    job1 = InferenceJob(name="job1", run=runner.job("job1", unblock))
    job2 = InferenceJob(name="job2", run=runner.job("job2"))
    job3 = InferenceJob(name="job3", run=runner.job("job3"))

    queue.submit(job1)
    queue.submit(job2)
    queue.submit(job3)

    import time
    time.sleep(0.01)

    queue.stop()

    # job1 was in-flight, so let it finish
    unblock.set()
    time.sleep(0.05)

    # Only job1 should have run; job2 and job3 were discarded
    names = [name for _, name in runner.events if _ == "end"]
    assert names == ["job1"]


def test_depth_reflects_queued_count():
    """depth() returns the number of queued (not in-flight) jobs."""
    queue = InferenceQueue(size=8)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()

    job1 = InferenceJob(name="job1", run=runner.job("job1", unblock))
    job2 = InferenceJob(name="job2", run=runner.job("job2"))
    job3 = InferenceJob(name="job3", run=runner.job("job3"))

    queue.submit(job1)
    import time
    time.sleep(0.05)
    # Give worker time to dequeue job1
    while queue.depth() > 0 and not runner.in_flight():
        time.sleep(0.001)

    queue.submit(job2)
    queue.submit(job3)
    # job1 is in-flight, job2 and job3 are queued
    assert queue.depth() >= 1

    unblock.set()
    time.sleep(0.1)

    # All jobs should finish; no more queued
    assert queue.depth() == 0

    queue.stop()


def test_no_submit_after_stop():
    """submit() returns False after stop()."""
    queue = InferenceQueue(size=4)
    queue.start()
    queue.stop()

    job = InferenceJob(name="job", run=lambda: None)
    assert queue.submit(job) is False


def test_bound_holds_when_only_priority_jobs_are_present():
    """A normal job is rejected when there is no normal job to evict."""
    queue = InferenceQueue(size=1)
    release = threading.Event()

    def block():
        release.wait()

    queue.start()
    try:
        assert queue.submit(InferenceJob(name="cmd", run=block, priority=-1)) is True
        assert queue.submit(InferenceJob(name="chat", run=lambda: None)) is False
    finally:
        release.set()
        queue.stop()


def test_in_flight_true_during_job_execution():
    """in_flight() returns True while a job is executing, False after."""
    queue = InferenceQueue(size=4)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()

    job1 = InferenceJob(name="job1", run=runner.job("job1", unblock))
    queue.submit(job1)

    import time
    time.sleep(0.01)

    # job1 is in-flight
    assert queue.in_flight() is True
    assert queue.depth() == 0

    unblock.set()
    time.sleep(0.05)

    # job1 finished
    assert queue.in_flight() is False

    queue.stop()


def test_submit_if_room_refuses_when_full_without_eviction():
    """1 in-flight + 3 queued at size=4 (reserve=1) refuses and mutates nothing."""
    queue = InferenceQueue(size=4)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()
    try:
        job1 = InferenceJob(name="in-flight", run=runner.job("in-flight", unblock))
        assert queue.submit(job1) is True
        _wait_in_flight(runner)

        job2 = InferenceJob(name="job2", run=runner.job("job2"))
        job3 = InferenceJob(name="job3", run=runner.job("job3"))
        job4 = InferenceJob(name="job4", run=runner.job("job4"))
        assert queue.submit(job2) is True
        assert queue.submit(job3) is True
        assert queue.submit(job4) is True

        # 1 in-flight + 3 queued == size; reserve=1 caps admission at 3.
        before = [id(job) for job in queue._queue]
        assert len(before) == 3

        assert queue.submit_if_room(InferenceJob(name="triage-extract", run=lambda: None)) is False

        after = [id(job) for job in queue._queue]
        assert after == before
        assert queue.depth() == 3
    finally:
        unblock.set()
        queue.stop()


def test_submit_if_room_admits_and_runs_when_room():
    """With room, submit_if_room appends the job and the worker runs it."""
    queue = InferenceQueue(size=4)
    ran = threading.Event()

    queue.start()
    try:
        job = InferenceJob(name="triage-extract", run=ran.set)
        assert queue.submit_if_room(job) is True
        assert ran.wait(timeout=1.0) is True
    finally:
        queue.stop()

    assert queue.depth() == 0


def test_submit_if_room_after_stop_refuses_and_enqueues_nothing():
    """After stop(), submit_if_room refuses and leaves the queue empty."""
    queue = InferenceQueue(size=4)
    queue.start()
    queue.stop()

    job = InferenceJob(name="triage-extract", run=lambda: None)
    assert queue.submit_if_room(job) is False
    assert queue.depth() == 0


def test_submit_if_room_reserve_larger_than_size_always_refuses():
    """A reserve >= size refuses without crashing, even on an empty queue."""
    queue = InferenceQueue(size=2)
    queue.start()
    try:
        assert queue.submit_if_room(InferenceJob(name="j", run=lambda: None), reserve=5) is False
        assert queue.depth() == 0
    finally:
        queue.stop()


def test_submit_still_evicts_oldest_normal_while_submit_if_room_refuses():
    """submit() keeps its drop-oldest behavior where submit_if_room refuses."""
    queue = InferenceQueue(size=3)
    runner = _JobRunner()
    unblock = threading.Event()

    queue.start()
    try:
        job1 = InferenceJob(name="in-flight", run=runner.job("in-flight", unblock), priority=0)
        assert queue.submit(job1) is True
        _wait_in_flight(runner)

        job2 = InferenceJob(name="job2", run=runner.job("job2"), priority=0)
        job3 = InferenceJob(name="job3", run=runner.job("job3"), priority=0)
        assert queue.submit(job2) is True
        assert queue.submit(job3) is True

        # 1 in-flight + 2 queued == size; reserve=1 caps admission at 2.
        refused = InferenceJob(name="triage-extract", run=lambda: None)
        assert queue.submit_if_room(refused) is False
        assert all(job is not refused for job in queue._queue)

        # submit() still evicts the oldest normal queued job (job2).
        job4 = InferenceJob(name="job4", run=lambda: None)
        assert queue.submit(job4) is True
        remaining_names = [job.name for job in queue._queue]
        assert remaining_names == ["job3", "job4"]
    finally:
        unblock.set()
        queue.stop()


def test_submit_if_room_concurrent_with_draining_worker_loses_no_job():
    """Admission from a producer thread racing the worker never raises or loses accepted jobs."""
    queue = InferenceQueue(size=8)
    accepted: list[str] = []
    executed: set[str] = set()
    errors: list[Exception] = []
    guard = threading.Lock()

    def make_run(name: str):
        def run():
            with guard:
                executed.add(name)

        return run

    def producer():
        try:
            for i in range(500):
                name = f"job{i}"
                job = InferenceJob(name=name, run=make_run(name))
                if queue.submit_if_room(job):
                    with guard:
                        accepted.append(name)
        except Exception as exc:  # noqa: BLE001 - recorded for assertion
            errors.append(exc)

    queue.start()
    try:
        thread = threading.Thread(target=producer)
        thread.start()
        thread.join(timeout=5.0)
        assert not thread.is_alive()
        _wait_idle(queue)

        assert errors == []
        assert accepted
        assert set(accepted) <= executed
    finally:
        queue.stop()


def test_stop_during_in_flight_is_bounded_then_refuses_admission():
    """stop() returns within its join bound while a job blocks; admission stays refused."""
    queue = InferenceQueue(size=4)
    started = threading.Event()
    release = threading.Event()

    def block():
        started.set()
        release.wait()

    queue.start()
    assert queue.submit(InferenceJob(name="long", run=block)) is True
    assert started.wait(timeout=1.0) is True

    begin = time.monotonic()
    queue.stop()
    elapsed = time.monotonic() - begin
    assert elapsed < 1.5
    assert queue.submit_if_room(InferenceJob(name="triage-extract", run=lambda: None)) is False
    assert queue.depth() == 0

    release.set()