Python unittest structure

.github/workflows/ci.yaml

@@ -8,16 +8,19 @@ on:
     branches: [main, dev]
 
 jobs:
+  pytest:
+    uses: ./.github/workflows/pytest.yaml
+    secrets: inherit
   mypy:
     uses: ./.github/workflows/mypy.yaml
     secrets: inherit
-  tests:
-    uses: ./.github/workflows/tests.yaml
+  e2e-tests:
+    uses: ./.github/workflows/e2e-tests.yaml
     secrets: inherit
   ci:
     runs-on: ubuntu-latest
     name: Run CI
-    needs: [mypy, tests]
+    needs: [mypy, pytest, e2e-tests]
     steps:
       - name: 'Done'
         run: echo "Done"

.github/workflows/tests.yaml → .github/workflows/e2e-tests.yaml

@@ -1,4 +1,4 @@
-name: Tests
+name: E2ETests
 
 on: [workflow_call]
 

.github/workflows/pytest.yaml

@@ -0,0 +1,23 @@
+name: Pytest
+
+on: [workflow_call]
+
+jobs:
+  mypy:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ./backend
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+          cache: 'pip'
+      - name: Install Poetry
+        run: pip install poetry
+      - name: Install dependencies
+        run: poetry install --with tests --with mypy --with custom-data
+      - name: Run Pytest
+        run: poetry run pytest --cov=chainlit/

.gitignore

@@ -55,3 +55,6 @@ dist-ssr
 *.njsproj
 *.sln
 *.sw?
+
+.aider*
+.coverage

backend/pyproject.toml

@@ -58,6 +58,9 @@ numpy = [
 optional = true
 
 [tool.poetry.group.tests.dependencies]
+pytest = "^8.3.2"
+pytest-asyncio = "^0.23.8"
+pytest-cov = "^5.0.0"
 openai = "^1.11.1"
 langchain = "^0.1.5"
 llama-index = "^0.10.45"
@@ -111,6 +114,7 @@ module = [
 ]
 ignore_missing_imports = true
 
+
 [tool.poetry.group.custom-data]
 optional = true
 
@@ -125,3 +129,8 @@ azure-storage-file-datalake = "^12.14.0"
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
+
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "auto"

backend/tests/__init__.py

@@ -0,0 +1 @@
+

backend/tests/test_context.py

@@ -0,0 +1,80 @@
+from unittest.mock import Mock
+
+import pytest
+from chainlit.context import (
+    ChainlitContext,
+    ChainlitContextException,
+    get_context,
+    init_http_context,
+    init_ws_context,
+)
+from chainlit.emitter import BaseChainlitEmitter, ChainlitEmitter
+from chainlit.session import HTTPSession, WebsocketSession
+
+
+@pytest.fixture
+def mock_websocket_session():
+    return Mock(spec=WebsocketSession)
+
+
+@pytest.fixture
+def mock_http_session():
+    return Mock(spec=HTTPSession)
+
+
+@pytest.fixture
+def mock_emitter():
+    return Mock(spec=BaseChainlitEmitter)
+
+
+async def test_chainlit_context_init_with_websocket(
+    mock_websocket_session, mock_emitter
+):
+    context = ChainlitContext(mock_websocket_session, mock_emitter)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+    assert context.session == mock_websocket_session
+    assert context.active_steps == []
+
+
+async def test_chainlit_context_init_with_http(mock_http_session):
+    context = ChainlitContext(mock_http_session)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+    assert context.session == mock_http_session
+    assert context.active_steps == []
+
+
+async def test_init_ws_context(mock_websocket_session):
+    context = init_ws_context(mock_websocket_session)
+    assert isinstance(context, ChainlitContext)
+    assert context.session == mock_websocket_session
+    assert isinstance(context.emitter, ChainlitEmitter)
+
+
+async def test_init_http_context():
+    context = init_http_context()
+    assert isinstance(context, ChainlitContext)
+    assert isinstance(context.session, HTTPSession)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+
+
+async def test_get_context():
+    with pytest.raises(ChainlitContextException):
+        get_context()
+
+    init_http_context()  # Initialize a context
+    context = get_context()
+    assert isinstance(context, ChainlitContext)
+
+
+async def test_current_step_and_run():
+    context = init_http_context()
+    assert context.current_step is None
+    assert context.current_run is None
+
+    # Mock a step
+    mock_step = Mock()
+    mock_step.name = "on_chat_start"
+    context.active_steps.append(mock_step)
+
+    assert context.current_step == mock_step
+    assert context.current_run == mock_step

backend/tests/test_emitter.py

@@ -0,0 +1,138 @@
+import pytest
+from unittest.mock import AsyncMock, MagicMock
+from chainlit.emitter import ChainlitEmitter
+from chainlit.element import ElementDict
+from chainlit.step import StepDict
+
+
+@pytest.fixture
+def mock_session():
+    session = MagicMock()
+    session.emit = AsyncMock()
+    return session
+
+
+@pytest.fixture
+def emitter(mock_session):
+    return ChainlitEmitter(mock_session)
+
+
+async def test_send_element(emitter: ChainlitEmitter, mock_session: MagicMock) -> None:
+    element_dict: ElementDict = {
+        "id": "test_element",
+        "threadId": None,
+        "type": "text",
+        "chainlitKey": None,
+        "url": None,
+        "objectKey": None,
+        "name": "Test Element",
+        "display": "inline",
+        "size": None,
+        "language": None,
+        "page": None,
+        "autoPlay": None,
+        "playerConfig": None,
+        "forId": None,
+        "mime": None
+    }
+
+    await emitter.send_element(element_dict)
+
+    mock_session.emit.assert_called_once_with("element", element_dict)
+
+
+async def test_send_step(emitter: ChainlitEmitter, mock_session: MagicMock) -> None:
+    step_dict: StepDict = {
+        "id": "test_step",
+        "type": "user_message",
+        "name": "Test Step",
+        "output": "This is a test step",
+    }
+
+    await emitter.send_step(step_dict)
+
+    mock_session.emit.assert_called_once_with("new_message", step_dict)
+
+
+async def test_update_step(emitter: ChainlitEmitter, mock_session: MagicMock) -> None:
+    step_dict: StepDict = {
+        "id": "test_step",
+        "type": "assistant_message",
+        "name": "Updated Test Step",
+        "output": "This is an updated test step",
+    }
+
+    await emitter.update_step(step_dict)
+
+    mock_session.emit.assert_called_once_with("update_message", step_dict)
+
+
+async def test_delete_step(emitter: ChainlitEmitter, mock_session: MagicMock) -> None:
+    step_dict: StepDict = {
+        "id": "test_step",
+        "type": "system_message",
+        "name": "Deleted Test Step",
+        "output": "This step will be deleted",
+    }
+
+    await emitter.delete_step(step_dict)
+
+    mock_session.emit.assert_called_once_with("delete_message", step_dict)
+
+
+async def test_send_timeout(emitter, mock_session):
+    await emitter.send_timeout("ask_timeout")
+    mock_session.emit.assert_called_once_with("ask_timeout", {})
+
+
+async def test_clear(emitter, mock_session):
+    await emitter.clear("clear_ask")
+    mock_session.emit.assert_called_once_with("clear_ask", {})
+
+
+async def test_send_token(emitter: ChainlitEmitter, mock_session: MagicMock) -> None:
+    await emitter.send_token("test_id", "test_token", is_sequence=True, is_input=False)
+    mock_session.emit.assert_called_once_with(
+        "stream_token",
+        {"id": "test_id", "token": "test_token", "isSequence": True, "isInput": False},
+    )
+
+
+async def test_set_chat_settings(emitter, mock_session):
+    settings = {"key": "value"}
+    emitter.set_chat_settings(settings)
+    assert emitter.session.chat_settings == settings
+
+
+async def test_send_action_response(emitter, mock_session):
+    await emitter.send_action_response("test_id", True, "Success")
+    mock_session.emit.assert_called_once_with(
+        "action_response", {"id": "test_id", "status": True, "response": "Success"}
+    )
+
+
+async def test_update_token_count(emitter, mock_session):
+    count = 100
+    await emitter.update_token_count(count)
+    mock_session.emit.assert_called_once_with("token_usage", count)
+
+
+async def test_task_start(emitter, mock_session):
+    await emitter.task_start()
+    mock_session.emit.assert_called_once_with("task_start", {})
+
+
+async def test_task_end(emitter, mock_session):
+    await emitter.task_end()
+    mock_session.emit.assert_called_once_with("task_end", {})
+
+
+async def test_stream_start(emitter: ChainlitEmitter, mock_session: MagicMock) -> None:
+    step_dict: StepDict = {
+        "id": "test_stream",
+        "type": "run",
+        "name": "Test Stream",
+        "output": "This is a test stream",
+    }
+    await emitter.stream_start(step_dict)
+    mock_session.emit.assert_called_once_with("stream_start", step_dict)

backend/tests/test_user_session.py

@@ -0,0 +1,54 @@
+import pytest
+import pytest_asyncio
+from unittest.mock import Mock
+from contextlib import asynccontextmanager
+from chainlit.user_session import UserSession
+from chainlit.context import ChainlitContext, context_var
+from chainlit.session import WebsocketSession
+
+
+@asynccontextmanager
+async def create_chainlit_context():
+    mock_session = Mock(spec=WebsocketSession)
+    mock_session.id = "test_session_id"
+    mock_session.user_env = {"test_env": "value"}
+    mock_session.chat_settings = {}
+    mock_session.user = None
+    mock_session.chat_profile = None
+    mock_session.http_referer = None
+    mock_session.client_type = "webapp"
+    mock_session.languages = ["en"]
+
+    context = ChainlitContext(mock_session)
+    token = context_var.set(context)
+    try:
+        yield context
+    finally:
+        context_var.reset(token)
+
+
+@pytest_asyncio.fixture
+async def mock_chainlit_context():
+    return create_chainlit_context()
+
+
+@pytest.fixture
+def user_session():
+    return UserSession()
+
+
+async def test_user_session_set_get(mock_chainlit_context, user_session):
+    async with mock_chainlit_context as context:
+        # Test setting a value
+        user_session.set("test_key", "test_value")
+
+        # Test getting the value
+        assert user_session.get("test_key") == "test_value"
+
+        # Test getting a default value for a non-existent key
+        assert user_session.get("non_existent_key", "default") == "default"
+
+        # Test getting session-related values
+        assert user_session.get("id") == context.session.id
+        assert user_session.get("env") == context.session.user_env
+        assert user_session.get("languages") == context.session.languages

package.json

@@ -23,7 +23,7 @@
     "prepare": "husky install",
     "lintUi": "cd frontend && pnpm run lint",
     "formatUi": "cd frontend && pnpm run format",
-    "lintPython": "cd backend && poetry run mypy chainlit/",
+    "lintPython": "cd backend && poetry run mypy chainlit/ tests/",
     "formatPython": "black `git ls-files | grep '.py$'` && isort --profile=black .",
     "buildUi": "cd libs/react-client && pnpm run build && cd ../copilot && pnpm run build && cd ../../frontend && pnpm run build",
     "build": "pnpm run buildUi && (mkdir -p backend/chainlit/frontend && cp -R frontend/dist backend/chainlit/frontend) && (mkdir -p backend/chainlit/copilot && cp -R libs/copilot/dist backend/chainlit/copilot) && (cd backend && poetry build)"