Files
chainlit--chainlit/backend/tests/test_element.py
T
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
chore: import upstream snapshot with attribution
2026-07-13 12:48:47 +08:00

585 lines
21 KiB
Python

import json
import uuid
from unittest.mock import AsyncMock
import pytest
from chainlit.element import (
Audio,
CustomElement,
Dataframe,
Element,
ElementDict,
File,
Image,
Pdf,
Task,
TaskList,
TaskStatus,
Text,
Video,
)
@pytest.mark.asyncio
class TestElementBase:
"""Test suite for the base Element class."""
async def test_element_initialization_with_url(self, mock_chainlit_context):
"""Test Element initialization with URL."""
async with mock_chainlit_context:
element = File(name="test_file", url="https://example.com/file.pdf")
assert element.name == "test_file"
assert element.url == "https://example.com/file.pdf"
assert isinstance(element.id, str)
uuid.UUID(element.id) # Verify valid UUID
assert element.persisted is False
assert element.updatable is False
async def test_element_initialization_with_content(self, mock_chainlit_context):
"""Test Element initialization with content."""
async with mock_chainlit_context:
content = b"test content"
element = File(name="test_file", content=content)
assert element.name == "test_file"
assert element.content == content
assert element.url is None
assert element.path is None
async def test_element_initialization_with_path(self, mock_chainlit_context):
"""Test Element initialization with path."""
async with mock_chainlit_context:
element = File(name="test_file", path="/path/to/file.txt")
assert element.name == "test_file"
assert element.path == "/path/to/file.txt"
assert element.url is None
assert element.content is None
async def test_element_requires_url_path_or_content(self, mock_chainlit_context):
"""Test that Element raises error without url, path, or content."""
async with mock_chainlit_context:
with pytest.raises(ValueError, match="Must provide url, path or content"):
File(name="test_file")
async def test_element_to_dict(self, mock_chainlit_context):
"""Test Element serialization to dictionary."""
async with mock_chainlit_context as ctx:
element = File(
name="test_file",
url="https://example.com/file.pdf",
display="inline",
)
element_dict = element.to_dict()
assert element_dict["name"] == "test_file"
assert element_dict["url"] == "https://example.com/file.pdf"
assert element_dict["type"] == "file"
assert element_dict["id"] == element.id
assert element_dict["threadId"] == ctx.session.thread_id
assert element_dict["display"] == "inline"
async def test_element_send(self, mock_chainlit_context):
"""Test Element.send() method."""
async with mock_chainlit_context as ctx:
element = File(name="test_file", url="https://example.com/file.pdf")
await element.send(for_id="message_123")
assert element.for_id == "message_123"
ctx.emitter.send_element.assert_called_once()
async def test_element_remove(self, mock_chainlit_context):
"""Test Element.remove() method."""
async with mock_chainlit_context as ctx:
element = File(name="test_file", url="https://example.com/file.pdf")
await element.remove()
ctx.emitter.emit.assert_called_once_with(
"remove_element", {"id": element.id}
)
async def test_element_display_options(self, mock_chainlit_context):
"""Test Element display options."""
async with mock_chainlit_context:
element_inline = File(
name="test", url="https://example.com/file.pdf", display="inline"
)
element_side = File(
name="test", url="https://example.com/file.pdf", display="side"
)
element_page = File(
name="test", url="https://example.com/file.pdf", display="page"
)
assert element_inline.display == "inline"
assert element_side.display == "side"
assert element_page.display == "page"
async def test_element_from_dict_file(self, mock_chainlit_context):
"""Test Element.from_dict() for File type."""
async with mock_chainlit_context:
element_dict: ElementDict = {
"id": str(uuid.uuid4()),
"name": "test_file",
"type": "file",
"url": "https://example.com/file.pdf",
"display": "inline",
}
element = Element.from_dict(element_dict)
assert isinstance(element, File)
assert element.name == "test_file"
assert element.url == "https://example.com/file.pdf"
async def test_element_from_dict_image(self, mock_chainlit_context):
"""Test Element.from_dict() for Image type."""
async with mock_chainlit_context:
element_dict: ElementDict = {
"id": str(uuid.uuid4()),
"name": "test_image",
"type": "image",
"url": "https://example.com/image.png",
"display": "inline",
}
element = Element.from_dict(element_dict)
assert isinstance(element, Image)
assert element.name == "test_image"
assert element.type == "image"
async def test_element_infer_type_from_mime(self):
"""Test Element.infer_type_from_mime() method."""
assert Element.infer_type_from_mime("image/png") == "image"
assert Element.infer_type_from_mime("image/jpeg") == "image"
assert Element.infer_type_from_mime("application/pdf") == "pdf"
assert Element.infer_type_from_mime("audio/mp3") == "audio"
assert Element.infer_type_from_mime("video/mp4") == "video"
assert Element.infer_type_from_mime("text/plain") == "file"
assert Element.infer_type_from_mime("application/json") == "file"
@pytest.mark.asyncio
class TestImageElement:
"""Test suite for Image element."""
async def test_image_initialization(self, mock_chainlit_context):
"""Test Image element initialization."""
async with mock_chainlit_context:
image = Image(
name="test_image",
url="https://example.com/image.png",
size="large",
)
assert image.type == "image"
assert image.name == "test_image"
assert image.size == "large"
async def test_image_size_options(self, mock_chainlit_context):
"""Test Image size options."""
async with mock_chainlit_context:
small = Image(name="test", url="https://example.com/img.png", size="small")
medium = Image(
name="test", url="https://example.com/img.png", size="medium"
)
large = Image(name="test", url="https://example.com/img.png", size="large")
assert small.size == "small"
assert medium.size == "medium"
assert large.size == "large"
@pytest.mark.asyncio
class TestTextElement:
"""Test suite for Text element."""
async def test_text_initialization(self, mock_chainlit_context):
"""Test Text element initialization."""
async with mock_chainlit_context:
text = Text(name="test_text", content="Hello, World!", language="python")
assert text.type == "text"
assert text.name == "test_text"
assert text.content == "Hello, World!"
assert text.language == "python"
async def test_text_without_language(self, mock_chainlit_context):
"""Test Text element without language."""
async with mock_chainlit_context:
text = Text(name="test_text", content="Plain text")
assert text.language is None
@pytest.mark.asyncio
class TestPdfElement:
"""Test suite for Pdf element."""
async def test_pdf_initialization(self, mock_chainlit_context):
"""Test Pdf element initialization."""
async with mock_chainlit_context:
pdf = Pdf(name="test_pdf", url="https://example.com/document.pdf", page=5)
assert pdf.type == "pdf"
assert pdf.name == "test_pdf"
assert pdf.mime == "application/pdf"
assert pdf.page == 5
async def test_pdf_without_page(self, mock_chainlit_context):
"""Test Pdf element without page number."""
async with mock_chainlit_context:
pdf = Pdf(name="test_pdf", url="https://example.com/document.pdf")
assert pdf.page is None
@pytest.mark.asyncio
class TestAudioElement:
"""Test suite for Audio element."""
async def test_audio_initialization(self, mock_chainlit_context):
"""Test Audio element initialization."""
async with mock_chainlit_context:
audio = Audio(
name="test_audio",
url="https://example.com/audio.mp3",
auto_play=True,
)
assert audio.type == "audio"
assert audio.name == "test_audio"
assert audio.auto_play is True
async def test_audio_default_auto_play(self, mock_chainlit_context):
"""Test Audio element default auto_play."""
async with mock_chainlit_context:
audio = Audio(name="test_audio", url="https://example.com/audio.mp3")
assert audio.auto_play is False
@pytest.mark.asyncio
class TestVideoElement:
"""Test suite for Video element."""
async def test_video_initialization(self, mock_chainlit_context):
"""Test Video element initialization."""
async with mock_chainlit_context:
player_config = {"youtube": {"playerVars": {"showinfo": 1}}}
video = Video(
name="test_video",
url="https://example.com/video.mp4",
size="large",
player_config=player_config,
)
assert video.type == "video"
assert video.name == "test_video"
assert video.size == "large"
assert video.player_config == player_config
async def test_video_without_player_config(self, mock_chainlit_context):
"""Test Video element without player config."""
async with mock_chainlit_context:
video = Video(name="test_video", url="https://example.com/video.mp4")
assert video.player_config is None
@pytest.mark.asyncio
class TestFileElement:
"""Test suite for File element."""
async def test_file_initialization(self, mock_chainlit_context):
"""Test File element initialization."""
async with mock_chainlit_context:
file = File(name="test_file", url="https://example.com/file.txt")
assert file.type == "file"
assert file.name == "test_file"
async def test_file_with_content(self, mock_chainlit_context):
"""Test File element with content."""
async with mock_chainlit_context:
content = b"File content"
file = File(name="test_file", content=content)
assert file.content == content
@pytest.mark.asyncio
class TestTaskListElement:
"""Test suite for TaskList element."""
async def test_tasklist_initialization(self, mock_chainlit_context):
"""Test TaskList element initialization."""
async with mock_chainlit_context:
tasklist = TaskList(name="test_tasklist")
assert tasklist.type == "tasklist"
assert tasklist.name == "test_tasklist"
assert tasklist.tasks == []
assert tasklist.status == "Ready"
assert tasklist.updatable is True
async def test_tasklist_add_task(self, mock_chainlit_context):
"""Test adding tasks to TaskList."""
async with mock_chainlit_context:
tasklist = TaskList(name="test_tasklist")
task1 = Task(title="Task 1", status=TaskStatus.READY)
task2 = Task(title="Task 2", status=TaskStatus.RUNNING)
await tasklist.add_task(task1)
await tasklist.add_task(task2)
assert len(tasklist.tasks) == 2
assert tasklist.tasks[0].title == "Task 1"
assert tasklist.tasks[1].title == "Task 2"
async def test_tasklist_preprocess_content(self, mock_chainlit_context):
"""Test TaskList content preprocessing."""
async with mock_chainlit_context:
tasklist = TaskList(name="test_tasklist", status="In Progress")
task = Task(title="Test Task", status=TaskStatus.DONE)
await tasklist.add_task(task)
await tasklist.preprocess_content()
assert isinstance(tasklist.content, str)
assert "Test Task" in tasklist.content
assert "done" in tasklist.content
assert "In Progress" in tasklist.content
@pytest.mark.asyncio
class TestTaskClass:
"""Test suite for Task class."""
def test_task_initialization(self):
"""Test Task initialization."""
task = Task(title="Test Task", status=TaskStatus.READY)
assert task.title == "Test Task"
assert task.status == TaskStatus.READY
assert task.forId is None
def test_task_with_for_id(self):
"""Test Task with forId."""
task = Task(title="Test Task", status=TaskStatus.RUNNING, forId="step_123")
assert task.forId == "step_123"
def test_task_status_enum(self):
"""Test TaskStatus enum values."""
assert TaskStatus.READY.value == "ready"
assert TaskStatus.RUNNING.value == "running"
assert TaskStatus.FAILED.value == "failed"
assert TaskStatus.DONE.value == "done"
@pytest.mark.asyncio
class TestCustomElement:
"""Test suite for CustomElement."""
async def test_custom_element_initialization(self, mock_chainlit_context):
"""Test CustomElement initialization."""
async with mock_chainlit_context:
props = {"key1": "value1", "key2": 42}
custom = CustomElement(name="test_custom", props=props)
assert custom.type == "custom"
assert custom.name == "test_custom"
assert custom.props == props
assert custom.mime == "application/json"
assert custom.updatable is True
async def test_custom_element_content_serialization(self, mock_chainlit_context):
"""Test CustomElement content serialization."""
async with mock_chainlit_context:
props = {"nested": {"data": [1, 2, 3]}}
custom = CustomElement(name="test_custom", props=props)
assert isinstance(custom.content, str)
assert "nested" in custom.content
assert "data" in custom.content
async def test_custom_element_update(self, mock_chainlit_context):
"""Test CustomElement update method."""
async with mock_chainlit_context as ctx:
custom = CustomElement(
name="test_custom",
props={"key": "value"},
url="https://example.com/custom",
)
custom.for_id = "message_123"
await custom.update()
ctx.emitter.send_element.assert_called()
@pytest.mark.asyncio
class TestElementEdgeCases:
"""Test suite for Element edge cases."""
async def test_element_with_custom_id(self, mock_chainlit_context):
"""Test Element with custom ID."""
async with mock_chainlit_context:
custom_id = str(uuid.uuid4())
element = File(
id=custom_id, name="test_file", url="https://example.com/file.txt"
)
assert element.id == custom_id
async def test_element_with_object_key(self, mock_chainlit_context):
"""Test Element with object_key."""
async with mock_chainlit_context:
element = File(
name="test_file",
url="https://example.com/file.txt",
object_key="s3://bucket/key",
)
assert element.object_key == "s3://bucket/key"
async def test_element_with_chainlit_key(self, mock_chainlit_context):
"""Test Element with chainlit_key."""
async with mock_chainlit_context:
element = File(
name="test_file",
url="https://example.com/file.txt",
chainlit_key="chainlit_key_123",
)
assert element.chainlit_key == "chainlit_key_123"
async def test_element_send_without_url_or_key_raises_error(
self, mock_chainlit_context
):
"""Test that send() raises error without url or chainlit_key."""
async with mock_chainlit_context as ctx:
# Mock persist_file to not set chainlit_key
ctx.session.persist_file = AsyncMock(return_value={"id": None})
element = File(name="test_file", content=b"test content")
with pytest.raises(ValueError, match="Must provide url or chainlit key"):
await element.send(for_id="message_123", persist=False)
async def test_element_from_dict_with_missing_fields(self, mock_chainlit_context):
"""Test Element.from_dict() with minimal fields."""
async with mock_chainlit_context:
element_dict: ElementDict = {
"type": "file",
"url": "https://example.com/file.txt",
}
element = Element.from_dict(element_dict)
assert isinstance(element, File)
assert element.name == ""
assert element.url == "https://example.com/file.txt"
async def test_element_id_uniqueness(self, mock_chainlit_context):
"""Test that each Element gets a unique ID."""
async with mock_chainlit_context:
element1 = File(name="file1", url="https://example.com/file1.txt")
element2 = File(name="file2", url="https://example.com/file2.txt")
element3 = File(name="file3", url="https://example.com/file3.txt")
ids = {element1.id, element2.id, element3.id}
assert len(ids) == 3 # All unique
@pytest.mark.asyncio
class TestDataframeElement:
"""Test suite for Dataframe element."""
async def test_dataframe_with_pandas(self, mock_chainlit_context):
"""Test Dataframe element with a pandas DataFrame."""
import pandas as pd
async with mock_chainlit_context:
df = pd.DataFrame({"a": [4, 2, 0], "b": ["foo", "bar", "baz"]})
element = Dataframe(name="test_df", data=df)
assert element.type == "dataframe"
assert element.size == "large"
parsed = json.loads(element.content)
assert parsed["columns"] == ["a", "b"]
assert parsed["data"] == [[4, "foo"], [2, "bar"], [0, "baz"]]
assert parsed["index"] == [0, 1, 2]
async def test_dataframe_with_polars(self, mock_chainlit_context):
"""Test Dataframe element with a polars DataFrame."""
import polars as pl
async with mock_chainlit_context:
df = pl.DataFrame({"a": [4, 2, 0], "b": ["foo", "bar", "baz"]})
element = Dataframe(name="test_df", data=df)
assert element.type == "dataframe"
assert element.size == "large"
parsed = json.loads(element.content)
assert parsed["columns"] == ["a", "b"]
assert parsed["data"] == [[4, "foo"], [2, "bar"], [0, "baz"]]
assert parsed["index"] == [0, 1, 2]
async def test_dataframe_with_invalid_data(self, mock_chainlit_context):
"""Test Dataframe element rejects non-DataFrame data."""
async with mock_chainlit_context:
with pytest.raises(
TypeError,
match=r"data must be a pandas\.DataFrame or polars\.DataFrame",
):
Dataframe(name="test_df", data={"a": [1, 2]})
async def test_dataframe_polars_and_pandas_produce_equal_outputs(
self, mock_chainlit_context
):
"""Test that pandas and polars DataFrames produce the same JSON."""
import pandas as pd
import polars as pl
async with mock_chainlit_context:
pd_df = pd.DataFrame({"a": [4, 2, 0], "b": ["foo", "bar", "baz"]})
pl_df = pl.DataFrame({"a": [4, 2, 0], "b": ["foo", "bar", "baz"]})
pd_element = Dataframe(name="pd_df", data=pd_df)
pl_element = Dataframe(name="pl_df", data=pl_df)
pd_parsed = json.loads(pd_element.content)
pl_parsed = json.loads(pl_element.content)
assert pd_parsed["columns"] == pl_parsed["columns"]
assert pd_parsed["index"] == pl_parsed["index"]
assert pd_parsed["data"] == pl_parsed["data"]
async def test_dataframe_polars_with_dates(self, mock_chainlit_context):
"""Test Dataframe element with polars date columns serializes correctly."""
from datetime import date
import polars as pl
async with mock_chainlit_context:
df = pl.DataFrame(
{"date": [date(2026, 1, 1), date(2025, 12, 31)], "val": [1, 2]}
)
element = Dataframe(name="test_df", data=df)
parsed = json.loads(element.content)
assert parsed["columns"] == ["date", "val"]
assert len(parsed["data"]) == 2
assert parsed["data"][0][0] == "2026-01-01"
assert parsed["data"][1][0] == "2025-12-31"