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