File size: 4,471 Bytes
dbdbe46 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
"""Tests for queue management system."""
import pytest
import os
import json
from datetime import datetime
from src.core.queue_manager import QueueManager, QueueItem
@pytest.fixture
def queue_manager(tmp_path):
"""Create queue manager with temporary directory."""
queue_dir = tmp_path / "queue"
return QueueManager(str(queue_dir))
def test_queue_manager_init(queue_manager):
"""Test queue manager initialization."""
assert os.path.exists(queue_manager.queue_dir)
assert queue_manager.queue.empty()
assert len(queue_manager.active_evaluations) == 0
def test_add_request(queue_manager):
"""Test adding requests to queue."""
request_id = queue_manager.add_request("org/model", "main")
assert not queue_manager.queue.empty()
assert os.path.exists(os.path.join(queue_manager.queue_dir, "queue_state.json"))
# Verify persisted state
with open(os.path.join(queue_manager.queue_dir, "queue_state.json")) as f:
state = json.load(f)
assert len(state) == 1
assert state[0]["model_id"] == "org/model"
def test_get_next_request(queue_manager):
"""Test retrieving requests from queue."""
added_id = queue_manager.add_request("org/model", "main")
item = queue_manager.get_next_request()
assert item is not None
assert item.model_id == "org/model"
assert item.revision == "main"
assert item.request_id in queue_manager.active_evaluations
def test_mark_complete(queue_manager):
"""Test marking requests as complete."""
added_id = queue_manager.add_request("org/model", "main")
item = queue_manager.get_next_request()
queue_manager.mark_complete(item.request_id)
assert item.request_id not in queue_manager.active_evaluations
def test_queue_status(queue_manager):
"""Test queue status reporting."""
queue_manager.add_request("org/model1", "main")
queue_manager.add_request("org/model2", "main")
item = queue_manager.get_next_request()
status = queue_manager.get_queue_status()
assert status["queued"] == 1
assert status["active"] == 1
assert item.request_id in status["active_evaluations"]
def test_priority_ordering(queue_manager):
"""Test priority-based queue ordering."""
# Add requests with different priorities
queue_manager.add_request("org/model1", "main", priority=2)
queue_manager.add_request("org/model2", "main", priority=1) # Higher priority
queue_manager.add_request("org/model3", "main", priority=3)
# First request should be model2 (priority 1)
item = queue_manager.get_next_request()
assert item.model_id == "org/model2"
# Second should be model1 (priority 2)
item = queue_manager.get_next_request()
assert item.model_id == "org/model1"
# Third should be model3 (priority 3)
item = queue_manager.get_next_request()
assert item.model_id == "org/model3"
def test_queue_persistence(tmp_path):
"""Test queue state persistence across instances."""
queue_dir = str(tmp_path / "queue")
# Create first instance and add requests
manager1 = QueueManager(queue_dir)
manager1.add_request("org/model1", "main")
manager1.add_request("org/model2", "main")
# Create second instance and verify state loaded
manager2 = QueueManager(queue_dir)
assert manager2.queue.qsize() == 2
# Verify requests can be retrieved in correct order
item1 = manager2.get_next_request()
assert item1.model_id == "org/model1"
item2 = manager2.get_next_request()
assert item2.model_id == "org/model2"
def test_concurrent_access(queue_manager):
"""Test concurrent queue access."""
import threading
import time
def add_and_get():
# Add a request
queue_manager.add_request("org/model", "main")
time.sleep(0.1) # Simulate some work
# Try to get a request
item = queue_manager.get_next_request()
if item:
queue_manager.mark_complete(item.request_id)
# Create multiple threads
threads = []
for _ in range(5):
thread = threading.Thread(target=add_and_get)
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Verify queue state is consistent
status = queue_manager.get_queue_status()
assert len(status["active_evaluations"]) == 0 # All should be marked complete
|