privateGPT/tests/server/chat/test_chat_async_worker.py

from unittest.mock import AsyncMock, MagicMock

import pytest

from private_gpt.server.chat_async.chat_async_service import ChatAsyncService


def _service(
    stream_manager: MagicMock | None = None,
    chat_facade: MagicMock | None = None,
) -> ChatAsyncService:
    if stream_manager is None:
        stream_manager = MagicMock()
        stream_manager.stream_exists = AsyncMock(return_value=False)
        stream_manager.cancel_stream = AsyncMock(return_value=False)
    if chat_facade is None:
        chat_facade = MagicMock()

    return ChatAsyncService(stream_manager=stream_manager, chat_facade=chat_facade)


@pytest.mark.anyio
async def test_initiate_chat_stream_starts_engine_event_processing() -> None:
    stream_manager = MagicMock()
    stream_manager.stream_exists = AsyncMock(return_value=False)
    stream_manager.create_and_start_stream = AsyncMock(return_value="msg-1")
    chat_facade = MagicMock()

    async def events():
        if False:
            yield None

    event_generator = events()
    chat_facade.create_chat_event_generator = AsyncMock(return_value=event_generator)

    service = _service(stream_manager=stream_manager, chat_facade=chat_facade)
    request = MagicMock()
    request.messages = []
    request.thinking.enabled = False
    updated_request = MagicMock()
    updated_request.messages = []
    updated_request.thinking.enabled = False
    request.model_copy.return_value = updated_request
    correlation_id = await service.initiate_chat_stream(
        request=request,
        message_id="msg-1",
    )

    assert correlation_id == "msg-1"
    request.context.model_copy.assert_called_once_with(
        update={"correlation_id": "msg-1"}
    )
    chat_facade.create_chat_event_generator.assert_awaited_once_with(
        request=updated_request
    )
    stream_manager.create_and_start_stream.assert_awaited_once_with(
        event_handler=stream_manager.create_and_start_stream.await_args.kwargs[
            "event_handler"
        ],
        stream_type="chat_completion",
        event_generator=event_generator,
        correlation_id="msg-1",
        metadata={"message_count": 0, "thinking_enabled": False},
    )


@pytest.mark.anyio
async def test_initiate_chat_stream_rejects_duplicate_message_id() -> None:
    stream_manager = MagicMock()
    stream_manager.stream_exists = AsyncMock(return_value=True)

    service = _service(stream_manager=stream_manager)

    with pytest.raises(ValueError, match="already exists"):
        await service.initiate_chat_stream(request=MagicMock(), message_id="msg-2")


@pytest.mark.anyio
async def test_cancel_only_delegates_to_stream_manager() -> None:
    stream_manager = MagicMock()
    stream_manager.cancel_stream = AsyncMock(return_value=False)
    chat_facade = MagicMock()
    chat_facade.cancel = AsyncMock(return_value=True)

    service = _service(stream_manager=stream_manager, chat_facade=chat_facade)

    await service.cancel_stream("msg-3")

    stream_manager.cancel_stream.assert_awaited_once_with("msg-3")
    chat_facade.cancel.assert_not_awaited()