learningcircuit--local-deep-research/tests/database/test_queue_service.py

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