tests/test_queue.py

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