wehub-resource-sync
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
851 lines
28 KiB
Python
851 lines
28 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
from chainlit import config
|
|
from chainlit.callbacks import password_auth_callback
|
|
from chainlit.data.base import BaseDataLayer
|
|
from chainlit.types import ThreadDict
|
|
from chainlit.user import User
|
|
|
|
|
|
async def test_password_auth_callback(test_config: config.ChainlitConfig):
|
|
@password_auth_callback
|
|
async def auth_func(username: str, password: str) -> User | None:
|
|
if username == "testuser" and password == "testpass": # nosec B105
|
|
return User(identifier="testuser")
|
|
return None
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.password_auth_callback is not None
|
|
|
|
# Test the wrapped function
|
|
result = await test_config.code.password_auth_callback("testuser", "testpass")
|
|
assert isinstance(result, User)
|
|
assert result.identifier == "testuser"
|
|
|
|
# Test with incorrect credentials
|
|
result = await test_config.code.password_auth_callback("wronguser", "wrongpass")
|
|
assert result is None
|
|
|
|
|
|
async def test_header_auth_callback(test_config: config.ChainlitConfig):
|
|
from starlette.datastructures import Headers
|
|
|
|
from chainlit.callbacks import header_auth_callback
|
|
|
|
@header_auth_callback
|
|
async def auth_func(headers: Headers) -> User | None:
|
|
if headers.get("Authorization") == "Bearer valid_token":
|
|
return User(identifier="testuser")
|
|
return None
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.header_auth_callback is not None
|
|
|
|
# Test the wrapped function with valid header
|
|
valid_headers = Headers({"Authorization": "Bearer valid_token"})
|
|
result = await test_config.code.header_auth_callback(valid_headers)
|
|
assert isinstance(result, User)
|
|
assert result.identifier == "testuser"
|
|
|
|
# Test with invalid header
|
|
invalid_headers = Headers({"Authorization": "Bearer invalid_token"})
|
|
result = await test_config.code.header_auth_callback(invalid_headers)
|
|
assert result is None
|
|
|
|
# Test with missing header
|
|
missing_headers = Headers({})
|
|
result = await test_config.code.header_auth_callback(missing_headers)
|
|
assert result is None
|
|
|
|
|
|
async def test_oauth_callback(test_config: config.ChainlitConfig):
|
|
from unittest.mock import patch
|
|
|
|
from chainlit.callbacks import oauth_callback
|
|
from chainlit.user import User
|
|
|
|
# Mock the get_configured_oauth_providers function
|
|
with patch(
|
|
"chainlit.callbacks.get_configured_oauth_providers", return_value=["google"]
|
|
):
|
|
|
|
@oauth_callback
|
|
async def auth_func(
|
|
provider_id: str,
|
|
token: str,
|
|
raw_user_data: dict,
|
|
default_app_user: User,
|
|
id_token: str | None = None,
|
|
) -> User | None:
|
|
if provider_id == "google" and token == "valid_token": # nosec B105
|
|
return User(identifier="oauth_user")
|
|
return None
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.oauth_callback is not None
|
|
|
|
# Test the wrapped function with valid data
|
|
result = await test_config.code.oauth_callback(
|
|
"google", "valid_token", {}, User(identifier="default_user")
|
|
)
|
|
assert isinstance(result, User)
|
|
assert result.identifier == "oauth_user"
|
|
|
|
# Test with invalid data
|
|
result = await test_config.code.oauth_callback(
|
|
"facebook", "invalid_token", {}, User(identifier="default_user")
|
|
)
|
|
assert result is None
|
|
|
|
|
|
async def test_on_message(mock_chainlit_context, test_config: config.ChainlitConfig):
|
|
from chainlit.callbacks import on_message
|
|
from chainlit.message import Message
|
|
|
|
async with mock_chainlit_context:
|
|
message_received = None
|
|
|
|
@on_message
|
|
async def handle_message(message: Message):
|
|
nonlocal message_received
|
|
message_received = message
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.on_message is not None
|
|
|
|
# Create a test message
|
|
test_message = Message(content="Test message", author="User")
|
|
|
|
# Call the registered callback
|
|
await test_config.code.on_message(test_message)
|
|
|
|
# Check that the message was received by our handler
|
|
assert message_received is not None
|
|
assert message_received.content == "Test message"
|
|
assert message_received.author == "User"
|
|
|
|
|
|
async def test_on_stop(mock_chainlit_context, test_config: config.ChainlitConfig):
|
|
from chainlit.callbacks import on_stop
|
|
|
|
async with mock_chainlit_context:
|
|
stop_called = False
|
|
|
|
@on_stop
|
|
async def handle_stop():
|
|
nonlocal stop_called
|
|
stop_called = True
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.on_stop is not None
|
|
|
|
# Call the registered callback
|
|
await test_config.code.on_stop()
|
|
|
|
# Check that the stop_called flag was set
|
|
assert stop_called
|
|
|
|
|
|
async def test_action_callback(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.action import Action
|
|
from chainlit.callbacks import action_callback
|
|
|
|
async with mock_chainlit_context:
|
|
action_handled = False
|
|
|
|
@action_callback("test_action")
|
|
async def handle_action(action: Action):
|
|
nonlocal action_handled
|
|
action_handled = True
|
|
assert action.name == "test_action"
|
|
|
|
# Test that the callback is properly registered
|
|
assert "test_action" in test_config.code.action_callbacks
|
|
|
|
# Call the registered callback
|
|
test_action = Action(name="test_action", payload={"value": "test_value"})
|
|
await test_config.code.action_callbacks["test_action"](test_action)
|
|
|
|
# Check that the action_handled flag was set
|
|
assert action_handled
|
|
|
|
|
|
async def test_on_settings_update(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import on_settings_update
|
|
|
|
async with mock_chainlit_context:
|
|
settings_updated = False
|
|
|
|
@on_settings_update
|
|
async def handle_settings_update(settings: dict):
|
|
nonlocal settings_updated
|
|
settings_updated = True
|
|
assert settings == {"test_setting": "test_value"}
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.on_settings_update is not None
|
|
|
|
# Call the registered callback
|
|
await test_config.code.on_settings_update({"test_setting": "test_value"})
|
|
|
|
# Check that the settings_updated flag was set
|
|
assert settings_updated
|
|
|
|
|
|
async def test_author_rename(test_config: config.ChainlitConfig):
|
|
from chainlit.callbacks import author_rename
|
|
|
|
@author_rename
|
|
async def rename_author(author: str) -> str:
|
|
if author == "AI":
|
|
return "Assistant"
|
|
return author
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.author_rename is not None
|
|
|
|
# Call the registered callback
|
|
result = await test_config.code.author_rename("AI")
|
|
assert result == "Assistant"
|
|
|
|
result = await test_config.code.author_rename("Human")
|
|
assert result == "Human"
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.author_rename is not None
|
|
|
|
# Call the registered callback
|
|
result = await test_config.code.author_rename("AI")
|
|
assert result == "Assistant"
|
|
|
|
result = await test_config.code.author_rename("Human")
|
|
assert result == "Human"
|
|
|
|
|
|
async def test_on_app_startup(test_config: config.ChainlitConfig):
|
|
"""Test the on_app_startup callback registration and execution for sync and async functions."""
|
|
from chainlit.callbacks import on_app_startup
|
|
|
|
# Test with synchronous function
|
|
sync_startup_called = False
|
|
|
|
@on_app_startup
|
|
def sync_startup():
|
|
nonlocal sync_startup_called
|
|
sync_startup_called = True
|
|
|
|
assert test_config.code.on_app_startup is not None, (
|
|
"Sync startup callback not registered"
|
|
)
|
|
# Call the wrapped function (which might be async due to wrap_user_function)
|
|
result = test_config.code.on_app_startup()
|
|
if asyncio.iscoroutine(result):
|
|
await result
|
|
assert sync_startup_called, "Sync startup function was not called"
|
|
|
|
# Reset for async test
|
|
test_config.code.on_app_startup = None # Explicitly clear previous registration
|
|
|
|
# Test with asynchronous function
|
|
async_startup_called = False
|
|
|
|
@on_app_startup
|
|
async def async_startup():
|
|
nonlocal async_startup_called
|
|
await asyncio.sleep(0) # Simulate async work
|
|
async_startup_called = True
|
|
|
|
assert test_config.code.on_app_startup is not None, (
|
|
"Async startup callback not registered"
|
|
)
|
|
# Call the wrapped function (which should be async)
|
|
result = test_config.code.on_app_startup()
|
|
assert asyncio.iscoroutine(result), (
|
|
"Async startup function did not return a coroutine"
|
|
)
|
|
await result
|
|
assert async_startup_called, "Async startup function was not called"
|
|
|
|
|
|
async def test_on_app_shutdown(test_config: config.ChainlitConfig):
|
|
"""Test the on_app_shutdown callback registration and execution for sync and async functions."""
|
|
from chainlit.callbacks import on_app_shutdown
|
|
|
|
# Test with synchronous function
|
|
sync_shutdown_called = False
|
|
|
|
@on_app_shutdown
|
|
def sync_shutdown():
|
|
nonlocal sync_shutdown_called
|
|
sync_shutdown_called = True
|
|
|
|
assert test_config.code.on_app_shutdown is not None, (
|
|
"Sync shutdown callback not registered"
|
|
)
|
|
# Call the wrapped function
|
|
result = test_config.code.on_app_shutdown()
|
|
if asyncio.iscoroutine(result):
|
|
await result
|
|
assert sync_shutdown_called, "Sync shutdown function was not called"
|
|
|
|
# Reset for async test
|
|
test_config.code.on_app_shutdown = None # Explicitly clear previous registration
|
|
|
|
# Test with asynchronous function
|
|
async_shutdown_called = False
|
|
|
|
@on_app_shutdown
|
|
async def async_shutdown():
|
|
nonlocal async_shutdown_called
|
|
await asyncio.sleep(0) # Simulate async work
|
|
async_shutdown_called = True
|
|
|
|
assert test_config.code.on_app_shutdown is not None, (
|
|
"Async shutdown callback not registered"
|
|
)
|
|
# Call the wrapped function
|
|
result = test_config.code.on_app_shutdown()
|
|
assert asyncio.iscoroutine(result), (
|
|
"Async shutdown function did not return a coroutine"
|
|
)
|
|
await result
|
|
assert async_shutdown_called, "Async shutdown function was not called"
|
|
|
|
|
|
async def test_on_chat_start(mock_chainlit_context, test_config: config.ChainlitConfig):
|
|
from chainlit.callbacks import on_chat_start
|
|
|
|
async with mock_chainlit_context:
|
|
chat_started = False
|
|
|
|
@on_chat_start
|
|
async def handle_chat_start():
|
|
nonlocal chat_started
|
|
chat_started = True
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.on_chat_start is not None
|
|
|
|
# Call the registered callback
|
|
await test_config.code.on_chat_start()
|
|
|
|
# Check that the chat_started flag was set
|
|
assert chat_started
|
|
|
|
|
|
async def test_on_chat_resume(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import on_chat_resume
|
|
|
|
async with mock_chainlit_context:
|
|
chat_resumed = False
|
|
|
|
@on_chat_resume
|
|
async def handle_chat_resume(thread: ThreadDict):
|
|
nonlocal chat_resumed
|
|
chat_resumed = True
|
|
assert thread["id"] == "test_thread_id"
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.on_chat_resume is not None
|
|
|
|
# Call the registered callback
|
|
await test_config.code.on_chat_resume(
|
|
{
|
|
"id": "test_thread_id",
|
|
"createdAt": "2023-01-01T00:00:00Z",
|
|
"name": "Test Thread",
|
|
"userId": "test_user_id",
|
|
"userIdentifier": "test_user",
|
|
"tags": [],
|
|
"metadata": {},
|
|
"steps": [],
|
|
"elements": [],
|
|
}
|
|
)
|
|
|
|
# Check that the chat_resumed flag was set
|
|
assert chat_resumed
|
|
|
|
|
|
async def test_set_chat_profiles(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import set_chat_profiles
|
|
from chainlit.types import ChatProfile
|
|
|
|
async with mock_chainlit_context:
|
|
|
|
@set_chat_profiles
|
|
async def get_chat_profiles(user, language):
|
|
return [
|
|
ChatProfile(name="Test Profile", markdown_description="A test profile")
|
|
]
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.set_chat_profiles is not None
|
|
|
|
# Call the registered callback
|
|
result = await test_config.code.set_chat_profiles(None, None)
|
|
|
|
# Check the result
|
|
assert result is not None
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], ChatProfile)
|
|
assert result[0].name == "Test Profile"
|
|
assert result[0].markdown_description == "A test profile"
|
|
|
|
|
|
async def test_set_chat_profiles_language(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import set_chat_profiles
|
|
from chainlit.types import ChatProfile
|
|
|
|
async with mock_chainlit_context:
|
|
|
|
@set_chat_profiles
|
|
async def get_chat_profiles(user, language):
|
|
if language == "fr-CA":
|
|
return [
|
|
ChatProfile(
|
|
name="Profil de test", markdown_description="Un profil de test"
|
|
)
|
|
]
|
|
|
|
return [
|
|
ChatProfile(name="Test Profile", markdown_description="A test profile")
|
|
]
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.set_chat_profiles is not None
|
|
|
|
# Call the registered callback
|
|
result = await test_config.code.set_chat_profiles(None, "fr-CA")
|
|
|
|
# Check the result
|
|
assert result is not None
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], ChatProfile)
|
|
assert result[0].name == "Profil de test"
|
|
assert result[0].markdown_description == "Un profil de test"
|
|
|
|
|
|
async def test_set_starters(mock_chainlit_context, test_config: config.ChainlitConfig):
|
|
from chainlit.callbacks import set_starters
|
|
from chainlit.types import Starter
|
|
|
|
async with mock_chainlit_context:
|
|
|
|
@set_starters
|
|
async def get_starters(user):
|
|
return [
|
|
Starter(
|
|
label="Test Label",
|
|
message="Test Message",
|
|
)
|
|
]
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.set_starters is not None
|
|
|
|
# Call the registered callback
|
|
result = await test_config.code.set_starters(None, None)
|
|
|
|
# Check the result
|
|
assert result is not None
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], Starter)
|
|
assert result[0].label == "Test Label"
|
|
assert result[0].message == "Test Message"
|
|
|
|
|
|
async def test_set_starters_language(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import set_starters
|
|
from chainlit.types import Starter
|
|
|
|
async with mock_chainlit_context:
|
|
|
|
@set_starters
|
|
async def get_starters(user, language):
|
|
if language == "fr-CA":
|
|
return [
|
|
Starter(
|
|
label="Étiquette de test",
|
|
message="Message de test",
|
|
)
|
|
]
|
|
|
|
return [
|
|
Starter(
|
|
label="Test Label",
|
|
message="Test Message",
|
|
)
|
|
]
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.set_starters is not None
|
|
|
|
# Call the registered callback
|
|
result = await test_config.code.set_starters(None, "fr-CA")
|
|
|
|
# Check the result
|
|
assert result is not None
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], Starter)
|
|
assert result[0].label == "Étiquette de test"
|
|
assert result[0].message == "Message de test"
|
|
|
|
|
|
async def test_set_starter_categories(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import set_starter_categories
|
|
from chainlit.types import Starter, StarterCategory
|
|
|
|
async with mock_chainlit_context:
|
|
|
|
@set_starter_categories
|
|
async def get_starter_categories(user, language):
|
|
return [
|
|
StarterCategory(
|
|
label="Creative",
|
|
icon="https://example.com/creative.png",
|
|
starters=[
|
|
Starter(label="Write a poem", message="Write a poem"),
|
|
Starter(label="Write a story", message="Write a story"),
|
|
],
|
|
),
|
|
StarterCategory(
|
|
label="Educational",
|
|
starters=[
|
|
Starter(label="Explain concept", message="Explain it"),
|
|
],
|
|
),
|
|
]
|
|
|
|
assert test_config.code.set_starter_categories is not None
|
|
|
|
result = await test_config.code.set_starter_categories(None, None, None)
|
|
|
|
assert result is not None
|
|
assert isinstance(result, list)
|
|
assert len(result) == 2
|
|
|
|
assert result[0].label == "Creative"
|
|
assert result[0].icon == "https://example.com/creative.png"
|
|
assert len(result[0].starters) == 2
|
|
assert result[0].starters[0].label == "Write a poem"
|
|
|
|
assert result[1].label == "Educational"
|
|
assert result[1].icon is None
|
|
assert len(result[1].starters) == 1
|
|
|
|
category_dict = result[0].to_dict()
|
|
assert category_dict["label"] == "Creative"
|
|
assert category_dict["icon"] == "https://example.com/creative.png"
|
|
starters_list = category_dict["starters"]
|
|
assert isinstance(starters_list, list)
|
|
assert len(starters_list) == 2
|
|
|
|
|
|
async def test_set_starter_categories_with_chat_profile(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import set_starter_categories
|
|
from chainlit.types import Starter, StarterCategory
|
|
|
|
async with mock_chainlit_context:
|
|
|
|
@set_starter_categories
|
|
async def get_starter_categories(user, language, chat_profile):
|
|
return [
|
|
StarterCategory(
|
|
label="Profile Specific",
|
|
icon="https://example.com/profile.png",
|
|
starters=[
|
|
Starter(
|
|
label=f"Starter for {chat_profile}",
|
|
message=f"Message for {chat_profile}",
|
|
),
|
|
],
|
|
),
|
|
]
|
|
|
|
assert test_config.code.set_starter_categories is not None
|
|
|
|
result = await test_config.code.set_starter_categories(
|
|
None, "en", "test-profile"
|
|
)
|
|
|
|
assert result is not None
|
|
assert isinstance(result, list)
|
|
assert len(result) == 1
|
|
|
|
assert result[0].label == "Profile Specific"
|
|
assert result[0].icon == "https://example.com/profile.png"
|
|
assert len(result[0].starters) == 1
|
|
assert result[0].starters[0].label == "Starter for test-profile"
|
|
assert result[0].starters[0].message == "Message for test-profile"
|
|
|
|
|
|
async def test_on_shared_thread_view_allow(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import on_shared_thread_view
|
|
from chainlit.user import User
|
|
|
|
async with mock_chainlit_context:
|
|
# Simulate a viewer with access to certain chat profiles
|
|
allowed_profiles_by_user = {"viewer": {"pro", "basic"}}
|
|
|
|
@on_shared_thread_view
|
|
async def allow_shared_view(thread, viewer: User | None):
|
|
md = thread.get("metadata") or {}
|
|
chat_profile = (md or {}).get("chat_profile")
|
|
if not md.get("is_shared"):
|
|
return False
|
|
if not viewer:
|
|
return False
|
|
return chat_profile in allowed_profiles_by_user.get(
|
|
viewer.identifier, set()
|
|
)
|
|
|
|
assert test_config.code.on_shared_thread_view is not None
|
|
|
|
thread: ThreadDict = {
|
|
"id": "t1",
|
|
"createdAt": "2025-09-03T00:00:00Z",
|
|
"name": "Shared Thread",
|
|
"userId": "author_id",
|
|
"userIdentifier": "author",
|
|
"tags": [],
|
|
"metadata": {"is_shared": True, "chat_profile": "pro"},
|
|
"steps": [],
|
|
"elements": [],
|
|
}
|
|
viewer = User(identifier="viewer")
|
|
|
|
res = await test_config.code.on_shared_thread_view(thread, viewer)
|
|
assert res is True
|
|
|
|
|
|
async def test_on_shared_thread_view_block_and_exception(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
from chainlit.callbacks import on_shared_thread_view
|
|
from chainlit.user import User
|
|
|
|
async with mock_chainlit_context:
|
|
# Case 1: Explicitly return False when profile not allowed
|
|
@on_shared_thread_view
|
|
async def deny_when_not_allowed(thread, viewer: User | None):
|
|
md = thread.get("metadata") or {}
|
|
return md.get("chat_profile") == "allowed"
|
|
|
|
assert test_config.code.on_shared_thread_view is not None
|
|
|
|
thread: ThreadDict = {
|
|
"id": "t2",
|
|
"createdAt": "2025-09-03T00:00:00Z",
|
|
"name": "Shared Thread",
|
|
"userId": "author_id",
|
|
"userIdentifier": "author",
|
|
"tags": [],
|
|
"metadata": {"is_shared": True, "chat_profile": "restricted"},
|
|
"steps": [],
|
|
"elements": [],
|
|
}
|
|
viewer = User(identifier="viewer")
|
|
res = await test_config.code.on_shared_thread_view(thread, viewer)
|
|
assert not res
|
|
|
|
# Case 2: Raise an exception inside callback; wrapper should swallow and result should be falsy
|
|
@on_shared_thread_view
|
|
async def raise_on_forbidden(thread, viewer: User | None):
|
|
md = thread.get("metadata") or {}
|
|
if md.get("chat_profile") == "forbidden":
|
|
raise ValueError("Viewer not allowed for this profile")
|
|
return True
|
|
|
|
assert test_config.code.on_shared_thread_view is not None
|
|
|
|
thread_err: ThreadDict = {
|
|
"id": "t3",
|
|
"createdAt": "2025-09-03T00:00:00Z",
|
|
"name": "Shared Thread",
|
|
"userId": "author_id",
|
|
"userIdentifier": "author",
|
|
"tags": [],
|
|
"metadata": {"is_shared": True, "chat_profile": "forbidden"},
|
|
"steps": [],
|
|
"elements": [],
|
|
}
|
|
res2 = await test_config.code.on_shared_thread_view(thread_err, viewer)
|
|
assert not res2
|
|
|
|
|
|
async def test_on_chat_end(mock_chainlit_context, test_config: config.ChainlitConfig):
|
|
from chainlit.callbacks import on_chat_end
|
|
|
|
async with mock_chainlit_context:
|
|
chat_ended = False
|
|
|
|
@on_chat_end
|
|
async def handle_chat_end():
|
|
nonlocal chat_ended
|
|
chat_ended = True
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.on_chat_end is not None
|
|
|
|
# Call the registered callback
|
|
await test_config.code.on_chat_end()
|
|
|
|
# Check that the chat_ended flag was set
|
|
assert chat_ended
|
|
|
|
|
|
def test_data_layer_config(
|
|
mock_data_layer: AsyncMock,
|
|
test_config: config.ChainlitConfig,
|
|
mock_get_data_layer: Mock,
|
|
):
|
|
"""Test whether we can properly configure a data layer."""
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.data_layer is not None
|
|
|
|
# Call the registered callback
|
|
result = test_config.code.data_layer()
|
|
|
|
# Check that the result is an instance of MockDataLayer
|
|
assert isinstance(result, BaseDataLayer)
|
|
|
|
mock_get_data_layer.assert_called_once()
|
|
|
|
|
|
def test_chat_profile_with_config_overrides():
|
|
"""Test that ChatProfile can be created with config_overrides."""
|
|
from chainlit.config import (
|
|
ChainlitConfigOverrides,
|
|
FeaturesSettings,
|
|
McpFeature,
|
|
UISettings,
|
|
)
|
|
from chainlit.types import ChatProfile
|
|
|
|
# Test creating a profile without config_overrides
|
|
basic_profile = ChatProfile(
|
|
name="Basic Profile", markdown_description="A basic profile without overrides"
|
|
)
|
|
assert basic_profile.config_overrides is None
|
|
|
|
# Test creating a profile with config_overrides
|
|
config_overrides = ChainlitConfigOverrides(
|
|
features=FeaturesSettings(mcp=McpFeature(enabled=True)),
|
|
ui=UISettings(
|
|
name="Custom App Name",
|
|
description="Custom description",
|
|
default_theme="light",
|
|
),
|
|
)
|
|
|
|
profile_with_overrides = ChatProfile(
|
|
name="MCP Profile",
|
|
markdown_description="A profile with MCP enabled",
|
|
config_overrides=config_overrides,
|
|
)
|
|
|
|
# Verify the profile was created successfully
|
|
assert profile_with_overrides.name == "MCP Profile"
|
|
assert profile_with_overrides.config_overrides is not None
|
|
assert profile_with_overrides.config_overrides.features.mcp.enabled is True
|
|
assert profile_with_overrides.config_overrides.ui.name == "Custom App Name"
|
|
assert profile_with_overrides.config_overrides.ui.default_theme == "light"
|
|
|
|
|
|
async def test_set_chat_profiles_with_config_overrides(
|
|
mock_chainlit_context, test_config: config.ChainlitConfig
|
|
):
|
|
"""Test that set_chat_profiles callback works with profiles that have config_overrides."""
|
|
from chainlit.callbacks import set_chat_profiles
|
|
from chainlit.config import (
|
|
ChainlitConfigOverrides,
|
|
FeaturesSettings,
|
|
McpFeature,
|
|
UISettings,
|
|
)
|
|
from chainlit.types import ChatProfile
|
|
|
|
async with mock_chainlit_context:
|
|
|
|
@set_chat_profiles
|
|
async def get_chat_profiles(user, language):
|
|
return [
|
|
ChatProfile(
|
|
name="Basic Profile",
|
|
markdown_description="A basic profile without overrides",
|
|
),
|
|
ChatProfile(
|
|
name="MCP Profile",
|
|
markdown_description="A profile with MCP enabled",
|
|
config_overrides=ChainlitConfigOverrides(
|
|
features=FeaturesSettings(mcp=McpFeature(enabled=True)),
|
|
ui=UISettings(name="MCP Assistant", default_theme="dark"),
|
|
),
|
|
),
|
|
ChatProfile(
|
|
name="Light Theme Profile",
|
|
markdown_description="A profile with light theme",
|
|
config_overrides=ChainlitConfigOverrides(
|
|
ui=UISettings(name="Light Theme App", default_theme="light")
|
|
),
|
|
),
|
|
]
|
|
|
|
# Test that the callback is properly registered
|
|
assert test_config.code.set_chat_profiles is not None
|
|
|
|
# Call the registered callback
|
|
result = await test_config.code.set_chat_profiles(None, None)
|
|
|
|
# Check the result
|
|
assert result is not None
|
|
assert isinstance(result, list)
|
|
assert len(result) == 3
|
|
|
|
# Test basic profile
|
|
basic_profile = result[0]
|
|
assert basic_profile.name == "Basic Profile"
|
|
assert basic_profile.config_overrides is None
|
|
|
|
# Test MCP profile
|
|
mcp_profile = result[1]
|
|
assert mcp_profile.name == "MCP Profile"
|
|
assert mcp_profile.config_overrides is not None
|
|
assert mcp_profile.config_overrides.features.mcp.enabled is True
|
|
assert mcp_profile.config_overrides.ui.name == "MCP Assistant"
|
|
assert mcp_profile.config_overrides.ui.default_theme == "dark"
|
|
|
|
# Test light theme profile
|
|
light_profile = result[2]
|
|
assert light_profile.name == "Light Theme Profile"
|
|
assert light_profile.config_overrides is not None
|
|
assert light_profile.config_overrides.ui.name == "Light Theme App"
|
|
assert light_profile.config_overrides.ui.default_theme == "light"
|