def test_session_auth_token_generation(mock_session):
    """Test that Cortex correctly extracts the token from a session object."""
    cortex = Cortex(
        model="llama3.2-1b",
        user="test_user",
        account="test_account",
        session=mock_session,
    )

    # Test the token extraction
    assert cortex._generate_auth_header() == 'Snowflake Token="mock_jwt_token"'
Test that Cortex correctly extracts the token from a session object.
llama-index-integrations/llms/llama-index-llms-cortex/tests/test_integration_cortex.py