ForcePilot/backend/test/unit/channel/worker/test_pool.py

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock, MagicMock

import pytest

from yuxi.channel.worker.pool import WorkerPool, WorkerPoolConfig


@pytest.fixture
def mock_queue_port():
    port = MagicMock()
    port.ensure_group = AsyncMock()
    port.dequeue = AsyncMock(return_value=[])
    port.ack = AsyncMock()
    return port


@pytest.fixture
def worker_pool(mock_queue_port):
    config = WorkerPoolConfig(num_workers=2, max_concurrent=5, poll_timeout_ms=1000)
    return WorkerPool(
        queue_port=mock_queue_port,
        dispatch_fn=AsyncMock(),
        config=config,
        session_router=True,
    )


@pytest.fixture
def legacy_worker_pool(mock_queue_port):
    config = WorkerPoolConfig(num_workers=2, max_concurrent=5, poll_timeout_ms=1000)
    return WorkerPool(
        queue_port=mock_queue_port,
        dispatch_fn=AsyncMock(),
        config=config,
        session_router=False,
    )


@pytest.mark.unit
async def test_worker_pool_start_stop(worker_pool, mock_queue_port):
    await worker_pool.start()
    assert worker_pool.is_running is True
    assert len(worker_pool._tasks) == 3
    mock_queue_port.ensure_group.assert_awaited_once()

    await worker_pool.stop()
    assert worker_pool.is_running is False
    assert len(worker_pool._tasks) == 0


@pytest.mark.unit
async def test_worker_pool_start_without_session_router(legacy_worker_pool, mock_queue_port):
    await legacy_worker_pool.start()
    assert legacy_worker_pool.is_running is True
    assert len(legacy_worker_pool._tasks) == 2

    await legacy_worker_pool.stop()
    assert legacy_worker_pool.is_running is False


@pytest.mark.unit
async def test_worker_pool_start_without_ensure_group(mock_queue_port):
    del mock_queue_port.ensure_group
    config = WorkerPoolConfig(num_workers=1, max_concurrent=5, poll_timeout_ms=1000)
    pool = WorkerPool(
        queue_port=mock_queue_port,
        dispatch_fn=AsyncMock(),
        config=config,
        session_router=True,
    )
    await pool.start()
    assert pool.is_running is True
    await pool.stop()


@pytest.mark.unit
async def test_dispatch_loop_routes_messages(worker_pool, mock_queue_port):
    messages = [
        {"session_id": "sess-1", "content": "msg1", "_stream_id": "stream-1"},
        {"session_id": "sess-2", "content": "msg2", "_stream_id": "stream-2"},
    ]
    call_count = 0

    async def dequeue_side_effect(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return messages
        return []

    mock_queue_port.dequeue = AsyncMock(side_effect=dequeue_side_effect)

    await worker_pool.start()
    await asyncio.sleep(0.2)
    await worker_pool.stop()

    assert mock_queue_port.dequeue.await_count >= 1


@pytest.mark.unit
async def test_worker_loop_dispatches_and_acks(worker_pool, mock_queue_port):
    msg = {"session_id": "sess-1", "content": "hello", "_stream_id": "stream-1"}
    call_count = 0

    async def dequeue_side_effect(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return [msg]
        return []

    mock_queue_port.dequeue = AsyncMock(side_effect=dequeue_side_effect)

    await worker_pool.start()
    await asyncio.sleep(0.3)
    await worker_pool.stop()

    worker_pool._dispatch.assert_awaited()
    mock_queue_port.ack.assert_awaited_with("stream-1")


@pytest.mark.unit
async def test_worker_loop_handles_dispatch_error(worker_pool, mock_queue_port):
    worker_pool._dispatch = AsyncMock(side_effect=Exception("dispatch failed"))
    msg = {"session_id": "sess-1", "content": "hello", "_stream_id": "stream-1"}
    call_count = 0

    async def dequeue_side_effect(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return [msg]
        return []

    mock_queue_port.dequeue = AsyncMock(side_effect=dequeue_side_effect)

    await worker_pool.start()
    await asyncio.sleep(0.3)
    await worker_pool.stop()

    worker_pool._dispatch.assert_awaited()


@pytest.mark.unit
async def test_route_with_session_id(worker_pool):
    for i in range(worker_pool._config.num_workers):
        worker_pool._session_queues.append(asyncio.Queue(maxsize=worker_pool._config.max_concurrent))

    msg = {"session_id": "sess-abc", "content": "hello"}
    idx = hash("sess-abc") % worker_pool._config.num_workers
    queue = worker_pool._session_queues[idx]

    await worker_pool._route(msg)
    assert queue.qsize() == 1
    assert queue.get_nowait() == msg


@pytest.mark.unit
async def test_route_without_session_id(worker_pool):
    for i in range(worker_pool._config.num_workers):
        worker_pool._session_queues.append(asyncio.Queue(maxsize=worker_pool._config.max_concurrent))

    msg = {"content": "hello"}
    initial_rr = worker_pool._rr_index

    await worker_pool._route(msg)
    assert worker_pool._rr_index == initial_rr + 1
    assert worker_pool._session_queues[initial_rr % worker_pool._config.num_workers].qsize() == 1


@pytest.mark.unit
async def test_legacy_worker_loop(worker_pool, mock_queue_port):
    worker_pool._session_router = False
    msg = {"content": "legacy", "_stream_id": "stream-legacy"}
    call_count = 0

    async def dequeue_side_effect(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return [msg]
        return []

    mock_queue_port.dequeue = AsyncMock(side_effect=dequeue_side_effect)

    await worker_pool.start()
    await asyncio.sleep(0.3)
    await worker_pool.stop()

    worker_pool._dispatch.assert_awaited()
    mock_queue_port.ack.assert_awaited_with("stream-legacy")


@pytest.mark.unit
async def test_worker_loop_without_stream_id(worker_pool, mock_queue_port):
    msg = {"session_id": "sess-1", "content": "no stream"}
    call_count = 0

    async def dequeue_side_effect(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return [msg]
        return []

    mock_queue_port.dequeue = AsyncMock(side_effect=dequeue_side_effect)

    await worker_pool.start()
    await asyncio.sleep(0.3)
    await worker_pool.stop()

    worker_pool._dispatch.assert_awaited()
    mock_queue_port.ack.assert_not_awaited()


@pytest.mark.unit
async def test_dispatch_loop_error_handling(worker_pool, mock_queue_port):
    mock_queue_port.dequeue = AsyncMock(side_effect=Exception("dequeue error"))

    await worker_pool.start()
    await asyncio.sleep(0.3)
    await worker_pool.stop()

    assert mock_queue_port.dequeue.await_count >= 1