b7f52be4c9
CI / Run CI (push) Has been cancelled
CI / check-backend (push) Has been cancelled
CI / check-frontend (push) Has been cancelled
CI / tests (push) Has been cancelled
CI / e2e-tests (push) Has been cancelled
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
280 lines
9.3 KiB
Python
280 lines
9.3 KiB
Python
"""Tests for Modes system functionality."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
import chainlit as cl
|
|
from chainlit.emitter import ChainlitEmitter
|
|
from chainlit.mode import Mode, ModeOption
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_modes():
|
|
"""Fixture providing sample modes for testing."""
|
|
return [
|
|
Mode(
|
|
id="model",
|
|
name="Model",
|
|
options=[
|
|
ModeOption(
|
|
id="gpt-4",
|
|
name="GPT-4",
|
|
description="Most capable model",
|
|
icon="sparkles",
|
|
default=True,
|
|
),
|
|
ModeOption(
|
|
id="gpt-3.5-turbo",
|
|
name="GPT-3.5 Turbo",
|
|
description="Fast and efficient",
|
|
icon="bolt",
|
|
default=False,
|
|
),
|
|
],
|
|
),
|
|
Mode(
|
|
id="reasoning",
|
|
name="Reasoning",
|
|
options=[
|
|
ModeOption(id="high", name="High", description="Maximum depth"),
|
|
ModeOption(
|
|
id="medium", name="Medium", description="Balanced", default=True
|
|
),
|
|
ModeOption(id="low", name="Low", description="Quick responses"),
|
|
],
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestModeOption:
|
|
"""Test suite for ModeOption dataclass."""
|
|
|
|
def test_mode_option_required_fields(self):
|
|
"""Test ModeOption with required fields only."""
|
|
option = ModeOption(id="test", name="Test Option")
|
|
|
|
assert option.id == "test"
|
|
assert option.name == "Test Option"
|
|
assert option.description is None
|
|
assert option.icon is None
|
|
assert option.default is False
|
|
|
|
def test_mode_option_all_fields(self):
|
|
"""Test ModeOption with all fields."""
|
|
option = ModeOption(
|
|
id="gpt-4",
|
|
name="GPT-4",
|
|
description="Most capable model",
|
|
icon="sparkles",
|
|
default=True,
|
|
)
|
|
|
|
assert option.id == "gpt-4"
|
|
assert option.name == "GPT-4"
|
|
assert option.description == "Most capable model"
|
|
assert option.icon == "sparkles"
|
|
assert option.default is True
|
|
|
|
def test_mode_option_to_dict(self):
|
|
"""Test ModeOption serialization."""
|
|
option = ModeOption(
|
|
id="test",
|
|
name="Test",
|
|
description="Test desc",
|
|
icon="star",
|
|
default=True,
|
|
)
|
|
|
|
option_dict = option.to_dict()
|
|
|
|
assert option_dict["id"] == "test"
|
|
assert option_dict["name"] == "Test"
|
|
assert option_dict["description"] == "Test desc"
|
|
assert option_dict["icon"] == "star"
|
|
assert option_dict["default"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestMode:
|
|
"""Test suite for Mode dataclass."""
|
|
|
|
def test_mode_creation(self, mock_modes):
|
|
"""Test Mode creation with options."""
|
|
mode = mock_modes[0]
|
|
|
|
assert mode.id == "model"
|
|
assert mode.name == "Model"
|
|
assert len(mode.options) == 2
|
|
|
|
def test_mode_to_dict(self, mock_modes):
|
|
"""Test Mode serialization."""
|
|
mode = mock_modes[0]
|
|
mode_dict = mode.to_dict()
|
|
|
|
assert mode_dict["id"] == "model"
|
|
assert mode_dict["name"] == "Model"
|
|
assert len(mode_dict["options"]) == 2
|
|
assert mode_dict["options"][0]["id"] == "gpt-4"
|
|
|
|
def test_mode_default_option(self, mock_modes):
|
|
"""Test finding default option in mode."""
|
|
mode = mock_modes[0]
|
|
|
|
default_option = next(
|
|
(opt for opt in mode.options if opt.default), mode.options[0]
|
|
)
|
|
|
|
assert default_option.id == "gpt-4"
|
|
|
|
def test_mode_fallback_to_first(self, mock_modes):
|
|
"""Test fallback to first option when no default set."""
|
|
mode = Mode(
|
|
id="test",
|
|
name="Test",
|
|
options=[
|
|
ModeOption(id="opt1", name="Option 1"),
|
|
ModeOption(id="opt2", name="Option 2"),
|
|
],
|
|
)
|
|
|
|
default_option = next(
|
|
(opt for opt in mode.options if opt.default),
|
|
mode.options[0] if mode.options else None,
|
|
)
|
|
|
|
assert default_option is not None
|
|
assert default_option.id == "opt1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestMessageWithModes:
|
|
"""Test suite for Message with modes field."""
|
|
|
|
async def test_message_with_modes(self, mock_chainlit_context):
|
|
"""Test that Message can be created with modes field."""
|
|
async with mock_chainlit_context:
|
|
modes = {"model": "gpt-4", "reasoning": "high"}
|
|
message = cl.Message(content="Test message", modes=modes)
|
|
|
|
assert message.modes == modes
|
|
assert message.content == "Test message"
|
|
|
|
async def test_message_to_dict_includes_modes(self, mock_chainlit_context):
|
|
"""Test that Message.to_dict() includes the modes field."""
|
|
async with mock_chainlit_context:
|
|
modes = {"model": "gpt-4", "reasoning": "medium"}
|
|
message = cl.Message(content="Test", modes=modes)
|
|
message_dict = message.to_dict()
|
|
|
|
assert "modes" in message_dict
|
|
assert message_dict["modes"] == modes
|
|
|
|
async def test_message_from_dict_with_modes(self, mock_chainlit_context):
|
|
"""Test that Message.from_dict() correctly handles modes field."""
|
|
async with mock_chainlit_context:
|
|
message_dict = {
|
|
"id": "test-id",
|
|
"content": "Test message",
|
|
"modes": {"model": "gpt-3.5-turbo", "reasoning": "low"},
|
|
"type": "user_message",
|
|
"createdAt": "2024-01-01T00:00:00",
|
|
"output": "Test message",
|
|
}
|
|
message = cl.Message.from_dict(message_dict)
|
|
|
|
assert message.modes == {"model": "gpt-3.5-turbo", "reasoning": "low"}
|
|
assert message.content == "Test message"
|
|
|
|
async def test_message_without_modes(self, mock_chainlit_context):
|
|
"""Test that Message works without modes field (backward compatibility)."""
|
|
async with mock_chainlit_context:
|
|
message = cl.Message(content="Test message")
|
|
|
|
assert message.modes is None
|
|
message_dict = message.to_dict()
|
|
assert message_dict.get("modes") is None
|
|
|
|
async def test_message_send_with_modes(self, mock_chainlit_context):
|
|
"""Test that sending a message with modes works."""
|
|
async with mock_chainlit_context as ctx:
|
|
modes = {"model": "gpt-4", "reasoning": "high"}
|
|
message = cl.Message(content="Test", modes=modes)
|
|
|
|
with patch("chainlit.message.chat_context") as mock_chat_ctx:
|
|
with patch("chainlit.message.get_data_layer", return_value=None):
|
|
with patch("chainlit.message.config") as mock_config:
|
|
mock_config.code.author_rename = None
|
|
|
|
result = await message.send()
|
|
|
|
assert result == message
|
|
assert message.modes == modes
|
|
mock_chat_ctx.add.assert_called_once_with(message)
|
|
ctx.emitter.send_step.assert_called_once()
|
|
|
|
# Verify the dict sent to emitter includes modes
|
|
call_args = ctx.emitter.send_step.call_args[0][0]
|
|
assert call_args["modes"] == modes
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestEmitterSetModes:
|
|
"""Test suite for emitter set_modes functionality."""
|
|
|
|
async def test_set_modes(
|
|
self, mock_modes, mock_websocket_session: MagicMock
|
|
) -> None:
|
|
"""Test set_modes emits correct event."""
|
|
emitter = ChainlitEmitter(mock_websocket_session)
|
|
modes_dicts = [mode.to_dict() for mode in mock_modes]
|
|
|
|
await emitter.set_modes(mock_modes)
|
|
|
|
mock_websocket_session.emit.assert_called_once_with("set_modes", modes_dicts)
|
|
|
|
async def test_set_modes_empty_list(
|
|
self, mock_websocket_session: MagicMock
|
|
) -> None:
|
|
"""Test set_modes with empty list."""
|
|
emitter = ChainlitEmitter(mock_websocket_session)
|
|
|
|
await emitter.set_modes([])
|
|
|
|
mock_websocket_session.emit.assert_called_once_with("set_modes", [])
|
|
|
|
async def test_set_modes_single_mode(
|
|
self, mock_websocket_session: MagicMock
|
|
) -> None:
|
|
"""Test set_modes with single mode."""
|
|
emitter = ChainlitEmitter(mock_websocket_session)
|
|
mode = Mode(
|
|
id="model",
|
|
name="Model",
|
|
options=[ModeOption(id="gpt-4", name="GPT-4", default=True)],
|
|
)
|
|
|
|
await emitter.set_modes([mode])
|
|
|
|
mock_websocket_session.emit.assert_called_once()
|
|
call_args = mock_websocket_session.emit.call_args
|
|
assert call_args[0][0] == "set_modes"
|
|
assert len(call_args[0][1]) == 1
|
|
assert call_args[0][1][0]["id"] == "model"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestModeExports:
|
|
"""Test that Mode and ModeOption are properly exported."""
|
|
|
|
def test_mode_exported_from_chainlit(self):
|
|
"""Test Mode is exported from chainlit package."""
|
|
assert hasattr(cl, "Mode")
|
|
assert cl.Mode is Mode
|
|
|
|
def test_mode_option_exported_from_chainlit(self):
|
|
"""Test ModeOption is exported from chainlit package."""
|
|
assert hasattr(cl, "ModeOption")
|
|
assert cl.ModeOption is ModeOption
|