def test_complete_with_session_auth(mock_cortex_with_session):
    """Test the complete 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":" world"}}]}',
        ]
        mock_post.return_value = mock_response

        response = mock_cortex_with_session.complete("Test prompt")

        # Verify the response was processed correctly
        assert isinstance(response, CompletionResponse)
        assert response.text == "Hello world"

        # 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 complete method using session authentication.
llama-index-integrations/llms/llama-index-llms-cortex/tests/test_integration_cortex.py