chore: import upstream snapshot with attribution
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
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
This commit is contained in:
@@ -0,0 +1,850 @@
|
||||
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"
|
||||
Reference in New Issue
Block a user