STT-Assistant/tests/test_inference.py

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

import threading
from stt_daemon.assistant.inference import InferenceJob, InferenceQueue


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()