""" Tests for database/queue_service.py Tests cover: - UserQueueService initialization - Queue status management - Task metadata operations - Task status updates - Pending task retrieval - Task cleanup """ from unittest.mock import Mock from datetime import datetime, UTC class TestUserQueueServiceInit: """Tests for UserQueueService initialization.""" def test_init_with_session(self): """Test initialization with a session.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() service = UserQueueService(mock_session) assert service.session == mock_session class TestUpdateQueueStatus: """Tests for update_queue_status method.""" def test_updates_existing_status(self): """Test updating existing queue status.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_status = Mock() mock_session.query.return_value.first.return_value = mock_status service = UserQueueService(mock_session) service.update_queue_status(5, 10, "task-123") assert mock_status.active_tasks == 5 assert mock_status.queued_tasks == 10 assert mock_status.last_task_id == "task-123" mock_session.commit.assert_called_once() def test_creates_new_status_when_none_exists(self): """Test creating status when none exists.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_session.query.return_value.first.return_value = None service = UserQueueService(mock_session) service.update_queue_status(2, 5) mock_session.add.assert_called_once() mock_session.commit.assert_called_once() class TestGetQueueStatus: """Tests for get_queue_status method.""" def test_returns_status_dict(self): """Test returns status as dictionary.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_status = Mock() mock_status.active_tasks = 3 mock_status.queued_tasks = 7 mock_status.last_checked = datetime.now(UTC) mock_status.last_task_id = "task-456" mock_session.query.return_value.first.return_value = mock_status service = UserQueueService(mock_session) result = service.get_queue_status() assert result["active_tasks"] == 3 assert result["queued_tasks"] == 7 assert result["last_task_id"] == "task-456" def test_returns_none_when_no_status(self): """Test returns None when no status exists.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_session.query.return_value.first.return_value = None service = UserQueueService(mock_session) result = service.get_queue_status() assert result is None class TestAddTaskMetadata: """Tests for add_task_metadata method.""" def test_adds_task_metadata(self): """Test adding task metadata.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() # Mock for _increment_queue_count mock_status = Mock() mock_status.queued_tasks = 0 mock_session.query.return_value.first.return_value = mock_status service = UserQueueService(mock_session) service.add_task_metadata("task-1", "research", priority=5) mock_session.add.assert_called() mock_session.commit.assert_called() def test_increments_queue_count(self): """Test that queue count is incremented.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_status = Mock() mock_status.queued_tasks = 2 mock_session.query.return_value.first.return_value = mock_status service = UserQueueService(mock_session) service.add_task_metadata("task-1", "research") assert mock_status.queued_tasks == 3 class TestUpdateTaskStatus: """Tests for update_task_status method.""" def test_updates_task_to_processing(self): """Test updating task to processing status.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_task = Mock() mock_task.status = "queued" mock_status = Mock() mock_status.queued_tasks = 5 mock_status.active_tasks = 2 # Set up query chain mock_filter = Mock() mock_filter.first.return_value = mock_task mock_query = Mock() mock_query.filter_by.return_value = mock_filter mock_query.first.return_value = mock_status def query_side_effect(model): if hasattr(model, "task_id"): # TaskMetadata return mock_query return Mock(first=Mock(return_value=mock_status)) mock_session.query.side_effect = query_side_effect service = UserQueueService(mock_session) service.update_task_status("task-1", "processing") assert mock_task.status == "processing" def test_updates_task_to_completed(self): """Test updating task to completed status.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_task = Mock() mock_task.status = "processing" mock_status = Mock() mock_status.queued_tasks = 3 mock_status.active_tasks = 2 mock_filter = Mock() mock_filter.first.return_value = mock_task mock_query = Mock() mock_query.filter_by.return_value = mock_filter mock_query.first.return_value = mock_status def query_side_effect(model): if hasattr(model, "task_id"): return mock_query return Mock(first=Mock(return_value=mock_status)) mock_session.query.side_effect = query_side_effect service = UserQueueService(mock_session) service.update_task_status("task-1", "completed") assert mock_task.status == "completed" def test_updates_task_with_error(self): """Test updating task with error message.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_task = Mock() mock_task.status = "processing" mock_status = Mock() mock_status.queued_tasks = 0 mock_status.active_tasks = 1 mock_filter = Mock() mock_filter.first.return_value = mock_task mock_query = Mock() mock_query.filter_by.return_value = mock_filter mock_query.first.return_value = mock_status def query_side_effect(model): if hasattr(model, "task_id"): return mock_query return Mock(first=Mock(return_value=mock_status)) mock_session.query.side_effect = query_side_effect service = UserQueueService(mock_session) service.update_task_status("task-1", "failed", "Something went wrong") assert mock_task.status == "failed" assert mock_task.error_message == "Something went wrong" def test_handles_nonexistent_task(self): """Test handling when task doesn't exist.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_filter = Mock() mock_filter.first.return_value = None mock_query = Mock() mock_query.filter_by.return_value = mock_filter mock_session.query.return_value = mock_query service = UserQueueService(mock_session) # Should not raise service.update_task_status("nonexistent", "completed") class TestGetPendingTasks: """Tests for get_pending_tasks method.""" def test_returns_pending_tasks(self): """Test getting pending tasks.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_task1 = Mock() mock_task1.task_id = "task-1" mock_task1.task_type = "research" mock_task1.created_at = datetime.now(UTC) mock_task1.priority = 5 mock_task2 = Mock() mock_task2.task_id = "task-2" mock_task2.task_type = "analysis" mock_task2.created_at = datetime.now(UTC) mock_task2.priority = 3 mock_query = mock_session.query.return_value mock_query.filter_by.return_value.order_by.return_value.limit.return_value.all.return_value = [ mock_task1, mock_task2, ] service = UserQueueService(mock_session) result = service.get_pending_tasks(limit=5) assert len(result) == 2 assert result[0]["task_id"] == "task-1" assert result[0]["task_type"] == "research" assert result[1]["task_id"] == "task-2" def test_returns_empty_list_when_no_tasks(self): """Test returns empty list when no pending tasks.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_query = mock_session.query.return_value mock_query.filter_by.return_value.order_by.return_value.limit.return_value.all.return_value = [] service = UserQueueService(mock_session) result = service.get_pending_tasks() assert result == [] class TestCleanupOldTasks: """Tests for cleanup_old_tasks method.""" def test_deletes_old_tasks(self): """Test deleting old completed tasks.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_query = mock_session.query.return_value mock_query.filter.return_value.delete.return_value = 5 service = UserQueueService(mock_session) result = service.cleanup_old_tasks(days=7) assert result == 5 mock_session.commit.assert_called_once() class TestQueueCountHelpers: """Tests for queue count helper methods.""" def test_increment_queue_count_existing(self): """Test incrementing queue count with existing status.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_status = Mock() mock_status.queued_tasks = 5 mock_session.query.return_value.first.return_value = mock_status service = UserQueueService(mock_session) service._increment_queue_count() assert mock_status.queued_tasks == 6 def test_increment_queue_count_new_status(self): """Test incrementing queue count creates new status.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_session.query.return_value.first.return_value = None service = UserQueueService(mock_session) service._increment_queue_count() mock_session.add.assert_called_once() def test_update_queue_counts_clamps_to_zero(self): """Test that queue counts don't go below zero.""" from local_deep_research.database.queue_service import UserQueueService mock_session = Mock() mock_status = Mock() mock_status.queued_tasks = 1 mock_status.active_tasks = 1 mock_session.query.return_value.first.return_value = mock_status service = UserQueueService(mock_session) service._update_queue_counts(-5, -5) assert mock_status.queued_tasks == 0 assert mock_status.active_tasks == 0 class TestProcessingToQueuedRevert: """update_task_status('queued') after 'processing' reverts the counter. Regression: the queued-dispatch loop's SystemAtCapacityError branch re-queues a research without reverting the queued->processing counter claim, so each capacity-rejected retry leaked a slot into active_tasks and decremented queued_tasks. Under sustained capacity pressure queued_tasks drifts to 0, at which point _process_user_queue treats the queue as empty and stops dispatching the still-present rows. These use a real in-memory DB so the counter arithmetic is actually exercised. """ @staticmethod def _service(): from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from local_deep_research.database.models.queue import ( QueueStatus, TaskMetadata, ) from local_deep_research.database.queue_service import ( UserQueueService, ) engine = create_engine("sqlite:///:memory:") QueueStatus.__table__.create(engine) TaskMetadata.__table__.create(engine) session = sessionmaker(bind=engine)() return UserQueueService(session) def test_processing_then_queued_round_trips_counts(self): svc = self._service() svc.add_task_metadata("r1", "research") # queued 1 / active 0 assert svc.get_queue_status()["queued_tasks"] == 1 assert svc.get_queue_status()["active_tasks"] == 0 svc.update_task_status("r1", "processing") # queued 0 / active 1 s = svc.get_queue_status() assert s["queued_tasks"] == 0 assert s["active_tasks"] == 1 svc.update_task_status("r1", "queued") # reverted: queued 1 / active 0 s = svc.get_queue_status() assert s["queued_tasks"] == 1, ( "queued->processing claim must be reverted on re-queue" ) assert s["active_tasks"] == 0, ( "active slot must not leak after a capacity re-queue" ) def test_repeated_capacity_rejects_do_not_drift(self): """Three processing->queued cycles (simulating capacity retries) must leave the counter exactly where it started.""" svc = self._service() svc.add_task_metadata("r1", "research") for _ in range(3): svc.update_task_status("r1", "processing") svc.update_task_status("r1", "queued") s = svc.get_queue_status() assert s["queued_tasks"] == 1 assert s["active_tasks"] == 0