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
585 lines
21 KiB
Python
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"
|