""" Unit tests for the steering module. """ import asyncio from datetime import datetime from unittest.mock import AsyncMock, MagicMock import pytest from copilot.steering import ( ConversationManager, MessageQueue, Priority, QueuedMessage, QueueFullError, ShutdownSentinel, StreamingInputGenerator, SHUTDOWN_SENTINEL, ) class TestPriority: def test_priority_ordering(self): """Test that priority values are ordered correctly.""" assert Priority.LOW.value == 0 assert Priority.NORMAL.value == 1 assert Priority.HIGH.value == 2 assert Priority.URGENT.value == 3 assert Priority.URGENT > Priority.HIGH > Priority.NORMAL > Priority.LOW class TestQueuedMessage: def test_creation(self): """Test basic message creation.""" msg = QueuedMessage( request_id="test-1", content="Hello", priority=Priority.NORMAL, session_id="session-1", ) assert msg.request_id == "test-1" assert msg.content == "Hello" assert msg.priority == Priority.NORMAL assert msg.session_id == "session-1" assert isinstance(msg.queued_at, datetime) def test_priority_comparison(self): """Test that higher priority messages sort first.""" low = QueuedMessage( request_id="low", content="Low", priority=Priority.LOW, session_id="s", sequence_number=1, ) normal = QueuedMessage( request_id="normal", content="Normal", priority=Priority.NORMAL, session_id="s", sequence_number=2, ) urgent = QueuedMessage( request_id="urgent", content="Urgent", priority=Priority.URGENT, session_id="s", sequence_number=3, ) # Higher priority should sort first (be "less than") assert urgent < normal < low assert not low < urgent def test_fifo_within_same_priority(self): """Test FIFO ordering for messages with same priority.""" first = QueuedMessage( request_id="first", content="First", priority=Priority.NORMAL, session_id="s", sequence_number=1, ) second = QueuedMessage( request_id="second", content="Second", priority=Priority.NORMAL, session_id="s", sequence_number=2, ) # Earlier sequence should sort first assert first < second assert not second < first def test_comparison_with_sentinel(self): """Test that messages sort before shutdown sentinel.""" msg = QueuedMessage( request_id="msg", content="Test", priority=Priority.URGENT, session_id="s", ) assert msg < SHUTDOWN_SENTINEL assert not SHUTDOWN_SENTINEL < msg class TestShutdownSentinel: def test_singleton(self): """Test that SHUTDOWN_SENTINEL is used consistently.""" assert isinstance(SHUTDOWN_SENTINEL, ShutdownSentinel) def test_comparison_with_messages(self): """Test sentinel always sorts last.""" msg = QueuedMessage( request_id="msg", content="Test", priority=Priority.LOW, session_id="s", ) assert msg < SHUTDOWN_SENTINEL assert SHUTDOWN_SENTINEL > msg assert not SHUTDOWN_SENTINEL < msg def test_comparison_with_itself(self): """Test sentinel comparison with itself.""" sentinel1 = ShutdownSentinel() sentinel2 = ShutdownSentinel() assert not sentinel1 < sentinel2 assert sentinel1 <= sentinel2 assert sentinel1 >= sentinel2 class TestMessageQueue: @pytest.mark.asyncio async def test_put_and_get(self): """Test basic put and get operations.""" queue = MessageQueue(max_depth=10) msg = QueuedMessage( request_id="test", content="Hello", priority=Priority.NORMAL, session_id="session-1", ) await queue.put(msg) assert queue.qsize() == 1 result = await queue.get() assert result == msg assert queue.qsize() == 0 @pytest.mark.asyncio async def test_priority_ordering(self): """Test that messages are returned in priority order.""" queue = MessageQueue(max_depth=10) low = QueuedMessage( request_id="low", content="Low", priority=Priority.LOW, session_id="s", ) urgent = QueuedMessage( request_id="urgent", content="Urgent", priority=Priority.URGENT, session_id="s", ) normal = QueuedMessage( request_id="normal", content="Normal", priority=Priority.NORMAL, session_id="s", ) # Add in arbitrary order await queue.put(low) await queue.put(urgent) await queue.put(normal) # Should come out in priority order assert (await queue.get()).request_id == "urgent" assert (await queue.get()).request_id == "normal" assert (await queue.get()).request_id == "low" @pytest.mark.asyncio async def test_queue_full_error(self): """Test that QueueFullError is raised when queue is full.""" queue = MessageQueue(max_depth=2) msg1 = QueuedMessage( request_id="1", content="1", priority=Priority.NORMAL, session_id="s" ) msg2 = QueuedMessage( request_id="2", content="2", priority=Priority.NORMAL, session_id="s" ) msg3 = QueuedMessage( request_id="3", content="3", priority=Priority.NORMAL, session_id="s" ) await queue.put(msg1) await queue.put(msg2) with pytest.raises(QueueFullError): await queue.put(msg3) @pytest.mark.asyncio async def test_shutdown_signal(self): """Test shutdown signaling.""" queue = MessageQueue() assert not queue.is_shutdown() queue.signal_shutdown() assert queue.is_shutdown() # Should be able to get the sentinel result = await queue.get() assert isinstance(result, ShutdownSentinel) @pytest.mark.asyncio async def test_empty_and_full(self): """Test empty and full properties.""" queue = MessageQueue(max_depth=2) assert queue.empty() assert not queue.full() msg = QueuedMessage( request_id="1", content="1", priority=Priority.NORMAL, session_id="s" ) await queue.put(msg) assert not queue.empty() assert not queue.full() msg2 = QueuedMessage( request_id="2", content="2", priority=Priority.NORMAL, session_id="s" ) await queue.put(msg2) assert queue.full() class TestStreamingInputGenerator: @pytest.mark.asyncio async def test_yields_messages(self): """Test that generator yields messages in correct format.""" queue = MessageQueue() generator = StreamingInputGenerator(queue) msg = QueuedMessage( request_id="test-1", content="Hello world", priority=Priority.NORMAL, session_id="session-123", metadata={"custom": "value"}, ) await queue.put(msg) queue.signal_shutdown() messages = [] async for message in generator: messages.append(message) assert len(messages) == 1 assert messages[0]["type"] == "user" assert messages[0]["message"]["role"] == "user" assert messages[0]["message"]["content"] == "Hello world" assert messages[0]["metadata"]["request_id"] == "test-1" assert messages[0]["metadata"]["priority"] == "NORMAL" assert messages[0]["metadata"]["session_id"] == "session-123" assert messages[0]["metadata"]["custom"] == "value" @pytest.mark.asyncio async def test_stops_on_shutdown(self): """Test that generator stops when shutdown sentinel is received.""" queue = MessageQueue() generator = StreamingInputGenerator(queue) msg1 = QueuedMessage( request_id="1", content="First", priority=Priority.NORMAL, session_id="s" ) msg2 = QueuedMessage( request_id="2", content="Second", priority=Priority.NORMAL, session_id="s" ) await queue.put(msg1) await queue.put(msg2) queue.signal_shutdown() messages = [] async for message in generator: messages.append(message) assert len(messages) == 2 @pytest.mark.asyncio async def test_priority_order_preserved(self): """Test that messages come out in priority order.""" queue = MessageQueue() generator = StreamingInputGenerator(queue) low = QueuedMessage( request_id="low", content="Low", priority=Priority.LOW, session_id="s" ) urgent = QueuedMessage( request_id="urgent", content="Urgent", priority=Priority.URGENT, session_id="s", ) await queue.put(low) await queue.put(urgent) queue.signal_shutdown() messages = [] async for message in generator: messages.append(message) assert messages[0]["metadata"]["request_id"] == "urgent" assert messages[1]["metadata"]["request_id"] == "low" class TestConversationManager: @pytest.mark.asyncio async def test_queue_message(self): """Test queuing messages.""" mock_session = MagicMock() mock_session.session_id = "test-session" mock_session.send = AsyncMock(return_value="msg-id") manager = ConversationManager(mock_session) request_id = await manager.queue_message("Hello", Priority.NORMAL) assert request_id == "req-1" assert manager.is_started assert manager.queue_size >= 0 # May have been processed already await manager.stop() @pytest.mark.asyncio async def test_auto_start(self): """Test that manager auto-starts on first message.""" mock_session = MagicMock() mock_session.session_id = "test-session" mock_session.send = AsyncMock(return_value="msg-id") manager = ConversationManager(mock_session) assert not manager.is_started await manager.queue_message("Hello", Priority.NORMAL) assert manager.is_started await manager.stop() @pytest.mark.asyncio async def test_custom_request_id(self): """Test using custom request ID.""" mock_session = MagicMock() mock_session.session_id = "test-session" mock_session.send = AsyncMock(return_value="msg-id") manager = ConversationManager(mock_session) request_id = await manager.queue_message( "Hello", Priority.NORMAL, request_id="custom-123" ) assert request_id == "custom-123" await manager.stop() @pytest.mark.asyncio async def test_stop_without_start(self): """Test that stop() works even if never started.""" mock_session = MagicMock() mock_session.session_id = "test-session" manager = ConversationManager(mock_session) await manager.stop() # Should not raise @pytest.mark.asyncio async def test_context_manager(self): """Test async context manager usage.""" mock_session = MagicMock() mock_session.session_id = "test-session" mock_session.send = AsyncMock(return_value="msg-id") async with ConversationManager(mock_session) as manager: await manager.queue_message("Hello", Priority.NORMAL) assert manager.is_started # Should be stopped after exiting context assert not manager.is_started @pytest.mark.asyncio async def test_message_sends_to_session(self): """Test that queued messages are sent to session.""" mock_session = MagicMock() mock_session.session_id = "test-session" mock_session.send = AsyncMock(return_value="msg-id") manager = ConversationManager(mock_session) await manager.queue_message("Hello world", Priority.NORMAL) # Give the processor task time to run await asyncio.sleep(0.1) # Verify send was called mock_session.send.assert_called() call_args = mock_session.send.call_args[0][0] assert call_args["prompt"] == "Hello world" await manager.stop() @pytest.mark.asyncio async def test_priority_processing_order(self): """Test that messages are processed in priority order.""" mock_session = MagicMock() mock_session.session_id = "test-session" processed_prompts = [] async def capture_send(options): processed_prompts.append(options["prompt"]) return "msg-id" mock_session.send = capture_send manager = ConversationManager(mock_session) # Queue messages in arbitrary order await manager.queue_message("Low priority", Priority.LOW) await manager.queue_message("Urgent!", Priority.URGENT) await manager.queue_message("Normal", Priority.NORMAL) # Give processor time to work await asyncio.sleep(0.2) await manager.stop() # Should be processed in priority order assert processed_prompts == ["Urgent!", "Normal", "Low priority"] @pytest.mark.asyncio async def test_queue_full_error(self): """Test that QueueFullError is raised when queue is full.""" mock_session = MagicMock() mock_session.session_id = "test-session" # Use a blocking send to keep the queue from draining send_started = asyncio.Event() send_continue = asyncio.Event() async def blocking_send(options): send_started.set() await send_continue.wait() return "msg-id" mock_session.send = blocking_send manager = ConversationManager(mock_session, max_queue_depth=2) # Queue first message - this will be picked up by processor immediately await manager.queue_message("1", Priority.NORMAL) # Wait for send to start (message 1 is now being processed) await send_started.wait() # Now queue 2 more - these should fill the queue await manager.queue_message("2", Priority.NORMAL) await manager.queue_message("3", Priority.NORMAL) # This should fail since queue is full and send is blocked with pytest.raises(QueueFullError): await manager.queue_message("4", Priority.NORMAL) # Unblock and cleanup send_continue.set() await manager.stop(timeout=1.0) @pytest.mark.asyncio async def test_attachments_forwarded(self): """Test that attachments are forwarded to session.""" mock_session = MagicMock() mock_session.session_id = "test-session" mock_session.send = AsyncMock(return_value="msg-id") manager = ConversationManager(mock_session) attachments = [{"type": "file", "path": "/test/file.py"}] await manager.queue_message( "Check this file", Priority.NORMAL, attachments=attachments ) await asyncio.sleep(0.1) mock_session.send.assert_called() call_args = mock_session.send.call_args[0][0] assert call_args["attachments"] == attachments await manager.stop() class TestYieldWhileProcessing: """Tests to verify the SDK can accept yields during processing (spec requirement).""" @pytest.mark.asyncio async def test_yield_while_processing(self): """Verify that new messages can be queued while previous ones are processing.""" mock_session = MagicMock() mock_session.session_id = "test-session" first_processing_started = asyncio.Event() first_processing_done = asyncio.Event() messages_yielded = [] call_count = 0 async def slow_send(options): nonlocal call_count call_count += 1 messages_yielded.append(options["prompt"]) if call_count == 1: first_processing_started.set() # Wait for the test to queue more messages await first_processing_done.wait() return "msg-id" mock_session.send = slow_send manager = ConversationManager(mock_session) # Queue first message await manager.queue_message("message 1", Priority.NORMAL) # Wait for processing to start await first_processing_started.wait() # These yields must succeed while msg1 is still processing await manager.queue_message("message 2", Priority.NORMAL) await manager.queue_message("message 3", Priority.NORMAL) # Verify queue accepted the messages (non-blocking put) assert manager.queue_size >= 2 # Let first processing complete first_processing_done.set() # Give time for all messages to process await asyncio.sleep(0.3) await manager.stop() # All messages should have been processed assert messages_yielded == ["message 1", "message 2", "message 3"]