chore: import upstream snapshot with attribution

backend/tests/test_chat_context.py

@@ -0,0 +1,459 @@
+import asyncio
+from contextlib import contextmanager
+from unittest.mock import Mock, patch
+
+from chainlit.chat_context import chat_context, chat_contexts
+from chainlit.context import ChainlitContext, context_var
+
+
+@contextmanager
+def mock_chainlit_context(session=None):
+    """Context manager to set up and tear down Chainlit context."""
+    # Mock the event loop since we're not in an async context
+    mock_loop = Mock(spec=asyncio.AbstractEventLoop)
+
+    with patch("asyncio.get_running_loop", return_value=mock_loop):
+        mock_context = ChainlitContext(session=session)
+        token = context_var.set(mock_context)
+        try:
+            yield mock_context
+        finally:
+            context_var.reset(token)
+
+
+class TestChatContext:
+    """Test suite for ChatContext class."""
+
+    def setup_method(self):
+        """Clear chat_contexts before each test."""
+        chat_contexts.clear()
+
+    def teardown_method(self):
+        """Clear chat_contexts after each test."""
+        chat_contexts.clear()
+
+    def test_get_without_session(self):
+        """Test get returns empty list when no session exists."""
+        with mock_chainlit_context(session=None):
+            result = chat_context.get()
+            assert result == []
+
+    def test_get_with_new_session(self):
+        """Test get creates new chat context for new session."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.get()
+
+            assert result == []
+            assert "session_123" in chat_contexts
+            assert chat_contexts["session_123"] == []
+
+    def test_get_returns_copy(self):
+        """Test get returns a copy of the chat context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.get()
+
+            assert result == [mock_message]
+            # Verify it's a copy, not the original
+            assert result is not chat_contexts["session_123"]
+
+    def test_get_with_existing_messages(self):
+        """Test get returns existing messages."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1, mock_msg2]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.get()
+
+            assert len(result) == 2
+            assert mock_msg1 in result
+            assert mock_msg2 in result
+
+    def test_add_without_session(self):
+        """Test add does nothing when no session exists."""
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=None):
+            result = chat_context.add(mock_message)
+
+            assert result is None
+            assert len(chat_contexts) == 0
+
+    def test_add_with_new_session(self):
+        """Test add creates new chat context and adds message."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.add(mock_message)
+
+            assert result == mock_message
+            assert "session_123" in chat_contexts
+            assert mock_message in chat_contexts["session_123"]
+
+    def test_add_message_to_existing_context(self):
+        """Test add appends message to existing context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.add(mock_msg2)
+
+            assert result == mock_msg2
+            assert len(chat_contexts["session_123"]) == 2
+            assert mock_msg1 in chat_contexts["session_123"]
+            assert mock_msg2 in chat_contexts["session_123"]
+
+    def test_add_duplicate_message(self):
+        """Test add does not add duplicate messages."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.add(mock_message)
+
+            assert result == mock_message
+            assert len(chat_contexts["session_123"]) == 1
+
+    def test_remove_without_session(self):
+        """Test remove returns False when no session exists."""
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=None):
+            result = chat_context.remove(mock_message)
+
+            assert result is False
+
+    def test_remove_with_nonexistent_context(self):
+        """Test remove returns False when context doesn't exist."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.remove(mock_message)
+
+            assert result is False
+
+    def test_remove_nonexistent_message(self):
+        """Test remove returns False when message not in context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.remove(mock_msg2)
+
+            assert result is False
+            assert mock_msg1 in chat_contexts["session_123"]
+
+    def test_remove_existing_message(self):
+        """Test remove successfully removes message."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1, mock_msg2]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.remove(mock_msg1)
+
+            assert result is True
+            assert mock_msg1 not in chat_contexts["session_123"]
+            assert mock_msg2 in chat_contexts["session_123"]
+            assert len(chat_contexts["session_123"]) == 1
+
+    def test_clear_without_session(self):
+        """Test clear does nothing when no session exists."""
+        chat_contexts["session_123"] = [Mock()]
+
+        with mock_chainlit_context(session=None):
+            chat_context.clear()
+
+            # Original context should remain
+            assert "session_123" in chat_contexts
+
+    def test_clear_with_nonexistent_context(self):
+        """Test clear does nothing when context doesn't exist."""
+        mock_session = Mock()
+        mock_session.id = "session_456"
+
+        chat_contexts["session_123"] = [Mock()]
+
+        with mock_chainlit_context(session=mock_session):
+            chat_context.clear()
+
+            # Original context should remain
+            assert "session_123" in chat_contexts
+
+    def test_clear_existing_context(self):
+        """Test clear empties existing context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1, mock_msg2]
+
+        with mock_chainlit_context(session=mock_session):
+            chat_context.clear()
+
+            assert "session_123" in chat_contexts
+            assert chat_contexts["session_123"] == []
+
+    def test_to_openai_with_assistant_message(self):
+        """Test to_openai converts assistant messages correctly."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        mock_message.type = "assistant_message"
+        mock_message.content = "Hello, how can I help?"
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 1
+            assert result[0] == {
+                "role": "assistant",
+                "content": "Hello, how can I help?",
+            }
+
+    def test_to_openai_with_user_message(self):
+        """Test to_openai converts user messages correctly."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        mock_message.type = "user_message"
+        mock_message.content = "What is the weather?"
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 1
+            assert result[0] == {"role": "user", "content": "What is the weather?"}
+
+    def test_to_openai_with_system_message(self):
+        """Test to_openai converts system messages correctly."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        mock_message.type = "system_message"
+        mock_message.content = "You are a helpful assistant."
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 1
+            assert result[0] == {
+                "role": "system",
+                "content": "You are a helpful assistant.",
+            }
+
+    def test_to_openai_with_unknown_message_type(self):
+        """Test to_openai treats unknown types as system messages."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        mock_message.type = "unknown_type"
+        mock_message.content = "Unknown message"
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 1
+            assert result[0] == {"role": "system", "content": "Unknown message"}
+
+    def test_to_openai_with_multiple_messages(self):
+        """Test to_openai converts multiple messages correctly."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg1.type = "user_message"
+        mock_msg1.content = "Hello"
+
+        mock_msg2 = Mock()
+        mock_msg2.type = "assistant_message"
+        mock_msg2.content = "Hi there!"
+
+        mock_msg3 = Mock()
+        mock_msg3.type = "user_message"
+        mock_msg3.content = "How are you?"
+
+        chat_contexts["session_123"] = [mock_msg1, mock_msg2, mock_msg3]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 3
+            assert result[0] == {"role": "user", "content": "Hello"}
+            assert result[1] == {"role": "assistant", "content": "Hi there!"}
+            assert result[2] == {"role": "user", "content": "How are you?"}
+
+    def test_to_openai_with_empty_context(self):
+        """Test to_openai returns empty list for empty context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        chat_contexts["session_123"] = []
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert result == []
+
+    def test_to_openai_without_session(self):
+        """Test to_openai returns empty list when no session exists."""
+        with mock_chainlit_context(session=None):
+            result = chat_context.to_openai()
+
+            assert result == []
+
+
+class TestChatContextEdgeCases:
+    """Test suite for chat_context edge cases."""
+
+    def setup_method(self):
+        """Clear chat_contexts before each test."""
+        chat_contexts.clear()
+
+    def teardown_method(self):
+        """Clear chat_contexts after each test."""
+        chat_contexts.clear()
+
+    def test_multiple_sessions_isolated(self):
+        """Test that different sessions have isolated contexts."""
+        mock_session1 = Mock()
+        mock_session1.id = "session_1"
+
+        mock_session2 = Mock()
+        mock_session2.id = "session_2"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+
+        with mock_chainlit_context(session=mock_session1):
+            chat_context.add(mock_msg1)
+
+        with mock_chainlit_context(session=mock_session2):
+            chat_context.add(mock_msg2)
+
+        assert len(chat_contexts) == 2
+        assert mock_msg1 in chat_contexts["session_1"]
+        assert mock_msg2 in chat_contexts["session_2"]
+        assert mock_msg1 not in chat_contexts["session_2"]
+        assert mock_msg2 not in chat_contexts["session_1"]
+
+    def test_add_then_remove_then_add_again(self):
+        """Test adding, removing, and re-adding the same message."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            # Add
+            chat_context.add(mock_message)
+            assert len(chat_contexts["session_123"]) == 1
+
+            # Remove
+            result = chat_context.remove(mock_message)
+            assert result is True
+            assert len(chat_contexts["session_123"]) == 0
+
+            # Add again
+            chat_context.add(mock_message)
+            assert len(chat_contexts["session_123"]) == 1
+
+    def test_clear_then_add(self):
+        """Test adding messages after clearing context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            chat_context.add(mock_msg1)
+            chat_context.clear()
+            chat_context.add(mock_msg2)
+
+            result = chat_context.get()
+            assert len(result) == 1
+            assert mock_msg2 in result
+            assert mock_msg1 not in result
+
+    def test_to_openai_with_mixed_message_types(self):
+        """Test to_openai with various message types in sequence."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        messages = [
+            Mock(type="system_message", content="System prompt"),
+            Mock(type="user_message", content="User query"),
+            Mock(type="assistant_message", content="Assistant response"),
+            Mock(type="other_type", content="Other message"),
+        ]
+
+        chat_contexts["session_123"] = messages
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 4
+            assert result[0]["role"] == "system"
+            assert result[1]["role"] == "user"
+            assert result[2]["role"] == "assistant"
+            assert result[3]["role"] == "system"  # Unknown types default to system
+
+    def test_chat_context_singleton(self):
+        """Test that chat_context is a singleton instance."""
+        from chainlit.chat_context import chat_context as imported_context
+
+        assert chat_context is imported_context
+
+    def test_add_returns_message(self):
+        """Test that add returns the message for chaining."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.add(mock_message)
+
+            assert result is mock_message