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
130 lines
3.8 KiB
Python
130 lines
3.8 KiB
Python
import datetime
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from chainlit import config
|
|
from chainlit.callbacks import data_layer
|
|
from chainlit.context import ChainlitContext, context_var
|
|
from chainlit.data.base import BaseDataLayer
|
|
from chainlit.session import HTTPSession, WebsocketSession
|
|
from chainlit.user import PersistedUser
|
|
from chainlit.user_session import UserSession
|
|
|
|
|
|
@pytest.fixture
|
|
def persisted_test_user():
|
|
return PersistedUser(
|
|
id="test_user_id",
|
|
createdAt=datetime.datetime.now().isoformat(),
|
|
identifier="test_user_identifier",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_session_factory(persisted_test_user: PersistedUser) -> Callable[..., Mock]:
|
|
def create_mock_session(**kwargs) -> Mock:
|
|
mock = Mock(spec=WebsocketSession)
|
|
mock.user = kwargs.get("user", persisted_test_user)
|
|
mock.id = kwargs.get("id", "test_session_id")
|
|
mock.user_env = kwargs.get("user_env", {"test_env": "value"})
|
|
mock.chat_settings = kwargs.get("chat_settings", {})
|
|
mock.chat_profile = kwargs.get("chat_profile", None)
|
|
mock.environ = kwargs.get("environ", None)
|
|
mock.client_type = kwargs.get("client_type", "webapp")
|
|
mock.thread_id = kwargs.get("thread_id", "test_thread_id")
|
|
mock.emit = AsyncMock()
|
|
mock.has_first_interaction = kwargs.get("has_first_interaction", True)
|
|
mock.files = kwargs.get("files", {})
|
|
mock.files_spec = kwargs.get("files_spec", {})
|
|
|
|
return mock
|
|
|
|
return create_mock_session
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_session(mock_session_factory) -> Mock:
|
|
return mock_session_factory()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def create_chainlit_context(mock_session):
|
|
from chainlit.emitter import BaseChainlitEmitter
|
|
|
|
# Create a mock emitter with all necessary methods
|
|
mock_emitter = Mock(spec=BaseChainlitEmitter)
|
|
mock_emitter.send_step = AsyncMock()
|
|
mock_emitter.update_step = AsyncMock()
|
|
mock_emitter.delete_step = AsyncMock()
|
|
mock_emitter.stream_start = AsyncMock()
|
|
mock_emitter.send_element = AsyncMock()
|
|
mock_emitter.send_action = AsyncMock()
|
|
mock_emitter.remove_action = AsyncMock()
|
|
mock_emitter.emit = AsyncMock()
|
|
mock_emitter.set_chat_settings = Mock() # Sync method, not async
|
|
|
|
context = ChainlitContext(mock_session, emitter=mock_emitter)
|
|
token = context_var.set(context)
|
|
try:
|
|
yield context
|
|
finally:
|
|
context_var.reset(token)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def mock_chainlit_context(persisted_test_user, mock_session):
|
|
mock_session.user = persisted_test_user
|
|
return create_chainlit_context(mock_session)
|
|
|
|
|
|
@pytest.fixture
|
|
def user_session():
|
|
return UserSession()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_websocket_session():
|
|
session = Mock(spec=WebsocketSession)
|
|
session.emit = AsyncMock()
|
|
|
|
return session
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_http_session():
|
|
return Mock(spec=HTTPSession)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_data_layer(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
|
|
mock_data_layer = AsyncMock(spec=BaseDataLayer)
|
|
|
|
return mock_data_layer
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_get_data_layer(mock_data_layer: AsyncMock, test_config: config.ChainlitConfig):
|
|
# Instantiate mock data layer
|
|
mock_get_data_layer = Mock(return_value=mock_data_layer)
|
|
|
|
# Configure it using @data_layer decorator
|
|
return data_layer(mock_get_data_layer)
|
|
|
|
|
|
@pytest.fixture
|
|
def test_config(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
|
|
monkeypatch.setenv("CHAINLIT_ROOT_PATH", str(tmp_path))
|
|
|
|
test_config = config.load_config()
|
|
|
|
monkeypatch.setattr("chainlit.callbacks.config", test_config)
|
|
monkeypatch.setattr("chainlit.server.config", test_config)
|
|
monkeypatch.setattr("chainlit.config.config", test_config)
|
|
|
|
return test_config
|