privateGPT/tests/components/streaming/test_chat_scheduler.py

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest
from injector import Injector

from private_gpt.components.streaming.tasks.chat_scheduler import (
    ArqChatScheduler,
    ChatSchedulerFactory,
    LocalChatScheduler,
)


@pytest.fixture
def injector() -> Injector:
    return Injector()


@pytest.mark.anyio
async def test_arq_chat_scheduler_cancel_aborts_job(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    abort_chat_job = AsyncMock(return_value=True)
    monkeypatch.setattr(
        "private_gpt.components.streaming.tasks.chat_scheduler.abort_chat_job",
        abort_chat_job,
    )

    scheduler = ArqChatScheduler()
    cancelled = await scheduler.cancel("msg-arq-3")

    assert cancelled is True
    abort_chat_job.assert_awaited_once_with(correlation_id="msg-arq-3")


@pytest.mark.anyio
async def test_local_chat_scheduler_cancel_cancels_task() -> None:
    async def _work() -> None:
        await asyncio.sleep(100)

    task = asyncio.create_task(_work(), name="chat_msg-local-2")

    scheduler = LocalChatScheduler()
    cancelled = await scheduler.cancel("msg-local-2")

    assert cancelled is True
    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task


@pytest.mark.anyio
async def test_local_chat_scheduler_cancel_returns_false_when_no_task() -> None:
    scheduler = LocalChatScheduler()
    cancelled = await scheduler.cancel("nonexistent-msg")
    assert cancelled is False


def test_chat_scheduler_factory_selects_local_mode(injector: Injector) -> None:
    factory = ChatSchedulerFactory(
        settings=SimpleNamespace(
            scheduler=SimpleNamespace(chat=SimpleNamespace(mode="local"))
        ),
        injector=injector,
    )
    assert isinstance(factory.get(), LocalChatScheduler)


def test_chat_scheduler_factory_selects_arq_mode(injector: Injector) -> None:
    factory = ChatSchedulerFactory(
        settings=SimpleNamespace(
            scheduler=SimpleNamespace(chat=SimpleNamespace(mode="arq"))
        ),
        injector=injector,
    )
    assert isinstance(factory.get(), ArqChatScheduler)


def test_chat_scheduler_factory_raises_on_unknown_mode(injector: Injector) -> None:
    with pytest.raises(ValueError, match=r"Unknown scheduler\.chat\.mode"):
        ChatSchedulerFactory(
            settings=SimpleNamespace(
                scheduler=SimpleNamespace(chat=SimpleNamespace(mode="missing"))
            ),
            injector=injector,
        )