def test_chat_with_session_auth(mock_cortex_with_session):
    """Test the chat method using session authentication."""
    with patch("requests.post") as mock_post:
        # Configure the mock to return a properly formed response
        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'data: {"choices":[{"delta":{"content":"Hello"}}]}',
            b'data: {"choices":[{"delta":{"content":" from Cortex"}}]}',
        ]
        mock_post.return_value = mock_response

        messages = [ChatMessage(role=MessageRole.USER, content="Hi")]
        response = mock_cortex_with_session.chat(messages)

        # Verify the response was processed correctly
        assert isinstance(response, ChatResponse)
        assert response.message.content == "Hello from Cortex"
        assert response.message.role == MessageRole.ASSISTANT

        # Verify the request had the correct authorization header
        args, kwargs = mock_post.call_args
        headers = kwargs.get("headers", {})
        assert "Authorization" in headers
        assert headers["Authorization"] == 'Snowflake Token="mock_jwt_token"'
Test the chat method using session authentication.
llama-index-integrations/llms/llama-index-llms-cortex/tests/test_integration_cortex.py