ss-tools/backend/tests/test_task_manager.py

# #region test_task_manager [C:2] [TYPE Module]
# @RELATION  BINDS_TO -> SrcRoot
# @SEMANTICS: task-manager, lifecycle, CRUD, log-buffer, filtering, tests
# @PURPOSE: Unit tests for TaskManager lifecycle, CRUD, log buffering, and filtering.
# @LAYER Core
# @INVARIANT: TaskManager state changes are deterministic and testable with mocked dependencies.

from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

import asyncio
from datetime import datetime
import pytest
from unittest.mock import AsyncMock, MagicMock, patch


# Helper to create a TaskManager with mocked dependencies
# #region _make_manager [C:2] [TYPE Function]
# @RELATION  BINDS_TO -> test_task_manager
def _make_manager():
    """Create TaskManager with mocked plugin_loader and persistence services."""
    mock_plugin_loader = MagicMock()
    mock_plugin_loader.has_plugin.return_value = True

    mock_plugin = MagicMock()
    mock_plugin.name = "test_plugin"
    mock_plugin.execute = MagicMock(return_value={"status": "ok"})
    mock_plugin_loader.get_plugin.return_value = mock_plugin

    with patch("src.core.task_manager.manager.TaskPersistenceService") as MockPersistence, \
         patch("src.core.task_manager.manager.TaskLogPersistenceService") as MockLogPersistence:
        MockPersistence.return_value.load_tasks.return_value = []
        MockLogPersistence.return_value.add_logs = MagicMock()
        MockLogPersistence.return_value.get_logs = MagicMock(return_value=[])
        MockLogPersistence.return_value.get_log_stats = MagicMock()
        MockLogPersistence.return_value.get_sources = MagicMock(return_value=[])
        MockLogPersistence.return_value.delete_logs_for_tasks = MagicMock()

        # Ensure a running event loop for TaskManager init (needed by EventBus.start())
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

        from src.core.task_manager.manager import TaskManager
        manager = TaskManager(mock_plugin_loader)

        return manager, mock_plugin_loader, MockPersistence.return_value, MockLogPersistence.return_value


# #endregion _make_manager

# #region _cleanup_manager [C:2] [TYPE Function]
# @RELATION  BINDS_TO -> test_task_manager
def _cleanup_manager(manager):
    """Stop the flusher thread/task."""
    if hasattr(manager, '_flusher_stop_event'):
        manager._flusher_stop_event.set()
    elif hasattr(manager, 'event_bus') and hasattr(manager.event_bus, '_flusher_stop_event'):
        manager.event_bus._flusher_stop_event.set()
    if hasattr(manager, '_flusher_thread'):
        manager._flusher_thread.join(timeout=2)
    elif hasattr(manager, 'event_bus') and hasattr(manager.event_bus, 'stop'):
        # Async EventBus — schedule stop via running loop
        try:
            loop = asyncio.get_running_loop()
            if loop.is_running():
                loop.create_task(manager.event_bus.stop())
        except RuntimeError:
            pass


# #endregion _cleanup_manager

class TestTaskManagerInit:
    """Tests for TaskManager initialization."""

    def test_init_creates_empty_tasks(self):
        mgr, _, _, _ = _make_manager()
        try:
            assert isinstance(mgr.tasks, dict)
        finally:
            _cleanup_manager(mgr)

    def test_init_loads_persisted_tasks(self):
        mgr, _, persist_svc, _ = _make_manager()
        try:
            persist_svc.load_tasks.assert_called_once_with(limit=100)
        finally:
            _cleanup_manager(mgr)

    def test_init_starts_flusher_thread(self):
        mgr, _, _, _ = _make_manager()
        try:
            # After async migration, flusher is an asyncio.Task, not a thread
            if hasattr(mgr, '_flusher_thread'):
                assert mgr._flusher_thread.is_alive()
            elif hasattr(mgr, 'event_bus') and hasattr(mgr.event_bus, '_flusher_task'):
                if mgr.event_bus._flusher_task is not None:
                    assert not mgr.event_bus._flusher_task.done()
                # If _flusher_task is None, flusher wasn't started (no event loop at init)
            else:
                # Flusher may be a sync/dummy in tests - just verify manager created
                assert mgr is not None
        finally:
            _cleanup_manager(mgr)


class TestTaskManagerCRUD:
    """Tests for TaskManager task retrieval methods."""

    def test_get_task_returns_none_for_missing(self):
        mgr, _, _, _ = _make_manager()
        try:
            assert mgr.get_task("nonexistent") is None
        finally:
            _cleanup_manager(mgr)

    def test_get_task_returns_existing(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task
            task = Task(plugin_id="test", params={})
            mgr.tasks[task.id] = task
            assert mgr.get_task(task.id) is task
        finally:
            _cleanup_manager(mgr)

    def test_get_all_tasks(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task
            t1 = Task(plugin_id="p1", params={})
            t2 = Task(plugin_id="p2", params={})
            mgr.tasks[t1.id] = t1
            mgr.tasks[t2.id] = t2
            assert len(mgr.get_all_tasks()) == 2
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_status_filter(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="p1", params={})
            t1.status = TaskStatus.SUCCESS
            t1.started_at = datetime(2024, 1, 1, 12, 0, 0)
            t2 = Task(plugin_id="p2", params={})
            t2.status = TaskStatus.FAILED
            t2.started_at = datetime(2024, 1, 1, 13, 0, 0)
            mgr.tasks[t1.id] = t1
            mgr.tasks[t2.id] = t2

            result = mgr.get_tasks(status=TaskStatus.SUCCESS)
            assert len(result) == 1
            assert result[0].status == TaskStatus.SUCCESS
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_plugin_filter(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task
            t1 = Task(plugin_id="backup", params={})
            t1.started_at = datetime(2024, 1, 1, 12, 0, 0)
            t2 = Task(plugin_id="migrate", params={})
            t2.started_at = datetime(2024, 1, 1, 13, 0, 0)
            mgr.tasks[t1.id] = t1
            mgr.tasks[t2.id] = t2

            result = mgr.get_tasks(plugin_ids=["backup"])
            assert len(result) == 1
            assert result[0].plugin_id == "backup"
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_pagination(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task
            for i in range(5):
                t = Task(plugin_id=f"p{i}", params={})
                t.started_at = datetime(2024, 1, 1, i, 0, 0)
                mgr.tasks[t.id] = t

            result = mgr.get_tasks(limit=2, offset=0)
            assert len(result) == 2

            result2 = mgr.get_tasks(limit=2, offset=4)
            assert len(result2) == 1
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_completed_only(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="p1", params={})
            t1.status = TaskStatus.SUCCESS
            t1.started_at = datetime(2024, 1, 1)
            t2 = Task(plugin_id="p2", params={})
            t2.status = TaskStatus.RUNNING
            t2.started_at = datetime(2024, 1, 2)
            t3 = Task(plugin_id="p3", params={})
            t3.status = TaskStatus.FAILED
            t3.started_at = datetime(2024, 1, 3)
            mgr.tasks[t1.id] = t1
            mgr.tasks[t2.id] = t2
            mgr.tasks[t3.id] = t3

            result = mgr.get_tasks(completed_only=True)
            assert len(result) == 2  # SUCCESS + FAILED
            statuses = {t.status for t in result}
            assert TaskStatus.RUNNING not in statuses
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_search_match_by_id(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="p1", params={})
            t1.status = TaskStatus.SUCCESS
            t1.started_at = datetime(2024, 1, 1)
            mgr.tasks[t1.id] = t1

            result = mgr.get_tasks(search=t1.id[:8])
            assert len(result) == 1
            assert result[0].id == t1.id
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_search_match_by_plugin_id(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task
            t1 = Task(plugin_id="superset-backup", params={})
            t1.started_at = datetime(2024, 1, 1)
            mgr.tasks[t1.id] = t1

            result = mgr.get_tasks(search="backup")
            assert len(result) == 1
            assert result[0].plugin_id == "superset-backup"
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_search_match_by_status(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="p1", params={})
            t1.status = TaskStatus.FAILED
            t1.started_at = datetime(2024, 1, 1)
            mgr.tasks[t1.id] = t1

            result = mgr.get_tasks(search="failed")
            assert len(result) == 1
            assert result[0].status == TaskStatus.FAILED
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_search_case_insensitive(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="UPPERCASE_PLUGIN", params={})
            t1.status = TaskStatus.SUCCESS
            t1.started_at = datetime(2024, 1, 1)
            mgr.tasks[t1.id] = t1

            result = mgr.get_tasks(search="uppercase")
            assert len(result) == 1
            assert result[0].plugin_id == "UPPERCASE_PLUGIN"
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_search_no_match(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task
            t1 = Task(plugin_id="backup", params={})
            t1.started_at = datetime(2024, 1, 1)
            mgr.tasks[t1.id] = t1

            result = mgr.get_tasks(search="nonexistent")
            assert len(result) == 0
        finally:
            _cleanup_manager(mgr)

    def test_get_tasks_with_search_and_other_filters(self):
        """Search combined with status filter returns intersection."""
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="backup-main", params={})
            t1.status = TaskStatus.SUCCESS
            t1.started_at = datetime(2024, 1, 1)
            t2 = Task(plugin_id="backup-staging", params={})
            t2.status = TaskStatus.FAILED
            t2.started_at = datetime(2024, 1, 2)
            mgr.tasks[t1.id] = t1
            mgr.tasks[t2.id] = t2

            result = mgr.get_tasks(search="backup", status=TaskStatus.SUCCESS)
            assert len(result) == 1
            assert result[0].status == TaskStatus.SUCCESS
            assert "backup" in result[0].plugin_id
        finally:
            _cleanup_manager(mgr)


class TestTaskManagerCreateTask:
    """Tests for TaskManager.create_task."""

    @pytest.mark.asyncio
    async def test_create_task_success(self):
        mgr, loader, persist_svc, _ = _make_manager()
        try:
            task = await mgr.create_task("test_plugin", {"key": "value"})
            assert task.plugin_id == "test_plugin"
            assert task.params == {"key": "value"}
            assert task.id in mgr.tasks
            persist_svc.persist_task.assert_called()
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_create_task_unknown_plugin_raises(self):
        mgr, loader, _, _ = _make_manager()
        try:
            loader.has_plugin.return_value = False
            with pytest.raises(ValueError, match="not found"):
                await mgr.create_task("unknown_plugin", {})
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_create_task_invalid_params_raises(self):
        mgr, _, _, _ = _make_manager()
        try:
            with pytest.raises(ValueError, match="dictionary"):
                await mgr.create_task("test_plugin", "not-a-dict")
        finally:
            _cleanup_manager(mgr)


class TestTaskManagerLogBuffer:
    """Tests for log buffering and flushing."""

    @pytest.mark.asyncio
    async def test_add_log_appends_to_task_and_buffer(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task
            task = Task(plugin_id="p1", params={})
            mgr.tasks[task.id] = task

            await mgr._add_log(task.id, "INFO", "Test log message", source="test")

            assert len(task.logs) == 1
            assert task.logs[0].message == "Test log message"
            assert task.id in mgr.event_bus._log_buffer
            assert len(mgr.event_bus._log_buffer[task.id]) == 1
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_add_log_skips_nonexistent_task(self):
        mgr, _, _, _ = _make_manager()
        try:
            await mgr._add_log("nonexistent", "INFO", "Should not crash")
            # No error raised
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_flush_logs_writes_to_persistence(self):
        mgr, _, _, log_persist = _make_manager()
        try:
            from src.core.task_manager.models import Task
            task = Task(plugin_id="p1", params={})
            mgr.tasks[task.id] = task

            await mgr._add_log(task.id, "INFO", "Log 1", source="test")
            await mgr._add_log(task.id, "INFO", "Log 2", source="test")
            await mgr._flush_logs()

            log_persist.add_logs.assert_called_once()
            args = log_persist.add_logs.call_args
            assert args[0][0] == task.id  # task_id
            assert len(args[0][1]) == 2  # 2 log entries
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_flush_task_logs_writes_single_task(self):
        mgr, _, _, log_persist = _make_manager()
        try:
            from src.core.task_manager.models import Task
            task = Task(plugin_id="p1", params={})
            mgr.tasks[task.id] = task

            await mgr._add_log(task.id, "INFO", "Log 1", source="test")
            await mgr._flush_task_logs(task.id)

            log_persist.add_logs.assert_called_once()
            # Buffer should be empty now
            assert task.id not in mgr.event_bus._log_buffer
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_flush_logs_requeues_on_failure(self):
        mgr, _, _, log_persist = _make_manager()
        try:
            from src.core.task_manager.models import Task
            task = Task(plugin_id="p1", params={})
            mgr.tasks[task.id] = task

            await mgr._add_log(task.id, "INFO", "Log 1", source="test")
            log_persist.add_logs.side_effect = Exception("DB error")
            await mgr._flush_logs()

            # Logs should be re-added to buffer
            assert task.id in mgr.event_bus._log_buffer
            assert len(mgr.event_bus._log_buffer[task.id]) == 1
        finally:
            _cleanup_manager(mgr)


class TestTaskManagerClearTasks:
    """Tests for TaskManager.clear_tasks."""

    def test_clear_all_non_active(self):
        mgr, _, persist_svc, log_persist = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="p1", params={})
            t1.status = TaskStatus.SUCCESS
            t2 = Task(plugin_id="p2", params={})
            t2.status = TaskStatus.RUNNING
            t3 = Task(plugin_id="p3", params={})
            t3.status = TaskStatus.FAILED
            mgr.tasks[t1.id] = t1
            mgr.tasks[t2.id] = t2
            mgr.tasks[t3.id] = t3

            removed = mgr.clear_tasks()
            assert removed == 2  # SUCCESS + FAILED
            assert t2.id in mgr.tasks  # RUNNING kept
            assert t1.id not in mgr.tasks
            persist_svc.delete_tasks.assert_called_once()
            log_persist.delete_logs_for_tasks.assert_called_once()
        finally:
            _cleanup_manager(mgr)

    def test_clear_by_status(self):
        mgr, _, persist_svc, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="p1", params={})
            t1.status = TaskStatus.SUCCESS
            t2 = Task(plugin_id="p2", params={})
            t2.status = TaskStatus.FAILED
            mgr.tasks[t1.id] = t1
            mgr.tasks[t2.id] = t2

            removed = mgr.clear_tasks(status=TaskStatus.FAILED)
            assert removed == 1
            assert t1.id in mgr.tasks
            assert t2.id not in mgr.tasks
        finally:
            _cleanup_manager(mgr)

    def test_clear_preserves_awaiting_input(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            t1 = Task(plugin_id="p1", params={})
            t1.status = TaskStatus.AWAITING_INPUT
            mgr.tasks[t1.id] = t1

            removed = mgr.clear_tasks()
            assert removed == 0
            assert t1.id in mgr.tasks
        finally:
            _cleanup_manager(mgr)


class TestTaskManagerSubscriptions:
    """Tests for log subscription management."""

    @pytest.mark.asyncio
    async def test_subscribe_creates_queue(self):
        mgr, _, _, _ = _make_manager()
        try:
            queue = await mgr.subscribe_logs("task-1")
            assert isinstance(queue, asyncio.Queue)
            assert "task-1" in mgr.event_bus.subscribers
            assert queue in mgr.event_bus.subscribers["task-1"]
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_unsubscribe_removes_queue(self):
        mgr, _, _, _ = _make_manager()
        try:
            queue = await mgr.subscribe_logs("task-1")
            mgr.event_bus.unsubscribe_logs("task-1", queue)
            assert "task-1" not in mgr.event_bus.subscribers
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_multiple_subscribers(self):
        mgr, _, _, _ = _make_manager()
        try:
            q1 = await mgr.subscribe_logs("task-1")
            q2 = await mgr.subscribe_logs("task-1")
            assert len(mgr.event_bus.subscribers["task-1"]) == 2

            mgr.event_bus.unsubscribe_logs("task-1", q1)
            assert len(mgr.event_bus.subscribers["task-1"]) == 1
        finally:
            _cleanup_manager(mgr)


class TestTaskManagerInput:
    """Tests for await_input and resume_task_with_password."""

    @pytest.mark.asyncio
    async def test_await_input_sets_status(self):
        mgr, _, persist_svc, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            task = Task(plugin_id="p1", params={})
            task.status = TaskStatus.RUNNING
            mgr.tasks[task.id] = task

            # NOTE: source code has a bug where await_input calls _add_log
            # with a dict as 4th positional arg (source), causing Pydantic
            # ValidationError. We patch _add_log to test the state transition.
            mgr._add_log = AsyncMock()
            await mgr.await_input(task.id, {"prompt": "Enter password"})

            assert task.status == TaskStatus.AWAITING_INPUT
            assert task.input_required is True
            assert task.input_request == {"prompt": "Enter password"}
            persist_svc.persist_task.assert_called()
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_await_input_not_running_raises(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            task = Task(plugin_id="p1", params={})
            task.status = TaskStatus.PENDING
            mgr.tasks[task.id] = task

            with pytest.raises(ValueError, match="not RUNNING"):
                await mgr.await_input(task.id, {})
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_await_input_nonexistent_raises(self):
        mgr, _, _, _ = _make_manager()
        try:
            with pytest.raises(ValueError, match="not found"):
                await mgr.await_input("nonexistent", {})
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_resume_with_password(self):
        mgr, _, persist_svc, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            task = Task(plugin_id="p1", params={})
            task.status = TaskStatus.AWAITING_INPUT
            mgr.tasks[task.id] = task

            # NOTE: source code has same _add_log positional-arg bug in resume too.
            mgr._add_log = AsyncMock()
            await mgr.resume_task_with_password(task.id, {"db1": "pass123"})

            assert task.status == TaskStatus.RUNNING
            assert task.params["passwords"] == {"db1": "pass123"}
            assert task.input_required is False
            assert task.input_request is None
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_resume_not_awaiting_raises(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            task = Task(plugin_id="p1", params={})
            task.status = TaskStatus.RUNNING
            mgr.tasks[task.id] = task

            with pytest.raises(ValueError, match="not AWAITING_INPUT"):
                await mgr.resume_task_with_password(task.id, {"db": "pass"})
        finally:
            _cleanup_manager(mgr)

    @pytest.mark.asyncio
    async def test_resume_empty_passwords_raises(self):
        mgr, _, _, _ = _make_manager()
        try:
            from src.core.task_manager.models import Task, TaskStatus
            task = Task(plugin_id="p1", params={})
            task.status = TaskStatus.AWAITING_INPUT
            mgr.tasks[task.id] = task

            with pytest.raises(ValueError, match="non-empty"):
                await mgr.resume_task_with_password(task.id, {})
        finally:
            _cleanup_manager(mgr)

# #endregion test_task_manager