[async] Apply #2469; enhance OAUTH async tests

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_connection.py

@@ -61,6 +61,7 @@
     OAUTH_AUTHENTICATOR,
     OAUTH_AUTHORIZATION_CODE,
     OAUTH_CLIENT_CREDENTIALS,
+    PAT_WITH_EXTERNAL_SESSION,
     PROGRAMMATIC_ACCESS_TOKEN,
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
@@ -247,7 +248,7 @@ async def __open_connection(self):
                 self._validate_client_prefetch_threads()
             )
 
-        # Setup authenticator
+        # Setup authenticator - validation happens in __config
         auth = Auth(self.rest)
 
         if self._session_token and self._master_token:
@@ -380,6 +381,12 @@ async def __open_connection(self):
                 )
             elif self._authenticator == PROGRAMMATIC_ACCESS_TOKEN:
                 self.auth_class = AuthByPAT(self._token)
+            elif self._authenticator == PAT_WITH_EXTERNAL_SESSION:
+                # TODO: SNOW-2344581: add support for PAT with external session ID for async connection
+                raise ProgrammingError(
+                    msg="PAT with external session ID is not supported for async connection.",
+                    errno=ER_INVALID_VALUE,
+                )
             elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
                 self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
                     self._client_request_mfa_token if IS_LINUX else True

src/snowflake/connector/aio/_wif_util.py

@@ -88,16 +88,23 @@ async def create_gcp_attestation(
 
     If the application isn't running on GCP or no credentials were found, raises an error.
     """
-    res = await session_manager.request(
-        method="GET",
-        url=f"http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/identity?audience={SNOWFLAKE_AUDIENCE}",
-        headers={
-            "Metadata-Flavor": "Google",
-        },
-    )
+    try:
+        res = await session_manager.request(
+            method="GET",
+            url=f"http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/identity?audience={SNOWFLAKE_AUDIENCE}",
+            headers={
+                "Metadata-Flavor": "Google",
+            },
+        )
+
+        content = await res.content.read()
+        jwt_str = content.decode("utf-8")
+    except Exception as e:
+        raise ProgrammingError(
+            msg=f"Error fetching GCP metadata: {e}. Ensure the application is running on GCP.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
 
-    content = await res.content.read()
-    jwt_str = content.decode("utf-8")
     _, subject = extract_iss_and_sub_without_signature_verification(jwt_str)
     return WorkloadIdentityAttestation(
         AttestationProvider.GCP, jwt_str, {"sub": subject}
@@ -139,15 +146,22 @@ async def create_azure_attestation(
     if managed_identity_client_id:
         query_params += f"&client_id={managed_identity_client_id}"
 
-    res = await session_manager.request(
-        method="GET",
-        url=f"{url_without_query_string}?{query_params}",
-        headers=headers,
-    )
+    try:
+        res = await session_manager.request(
+            method="GET",
+            url=f"{url_without_query_string}?{query_params}",
+            headers=headers,
+        )
+
+        content = await res.content.read()
+        response_text = content.decode("utf-8")
+        response_data = json.loads(response_text)
+    except Exception as e:
+        raise ProgrammingError(
+            msg=f"Error fetching Azure metadata: {e}. Ensure the application is running on Azure.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
 
-    content = await res.content.read()
-    response_text = content.decode("utf-8")
-    response_data = json.loads(response_text)
     jwt_str = response_data.get("access_token")
     if not jwt_str:
         raise ProgrammingError(

test/unit/aio/test_auth_keypair_async.py

@@ -8,6 +8,7 @@
 from test.unit.aio.mock_utils import mock_connection
 from unittest.mock import Mock, PropertyMock, patch
 
+import pytest
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
@@ -34,7 +35,8 @@ async def _mock_auth_key_pair_rest_response(url, headers, body, **kwargs):
     return _mock_auth_key_pair_rest_response
 
 
-async def test_auth_keypair():
+@pytest.mark.parametrize("authenticator", ["SNOWFLAKE_JWT", "snowflake_jwt"])
+async def test_auth_keypair(authenticator):
     """Simple Key Pair test."""
     private_key_der, public_key_der_encoded = generate_key_pair(2048)
     application = "testapplication"
@@ -43,7 +45,7 @@ async def test_auth_keypair():
     auth_instance = AuthByKeyPair(private_key=private_key_der)
     auth_instance._retry_ctx.set_start_time()
     await auth_instance.handle_timeout(
-        authenticator="SNOWFLAKE_JWT",
+        authenticator=authenticator,
         service_name=None,
         account=account,
         user=user,

test/unit/aio/test_auth_mfa_async.py

@@ -4,10 +4,15 @@
 
 from unittest import mock
 
+import pytest
+
 from snowflake.connector.aio import SnowflakeConnection
 
 
-async def test_mfa_token_cache():
+@pytest.mark.parametrize(
+    "authenticator", ["USERNAME_PASSWORD_MFA", "username_password_mfa"]
+)
+async def test_mfa_token_cache(authenticator):
     with mock.patch(
         "snowflake.connector.aio._network.SnowflakeRestful.fetch",
     ):
@@ -18,7 +23,7 @@ async def test_mfa_token_cache():
                 account="account",
                 user="user",
                 password="password",
-                authenticator="username_password_mfa",
+                authenticator=authenticator,
                 client_store_temporary_credential=True,
                 client_request_mfa_token=True,
             ):
@@ -44,7 +49,7 @@ async def test_mfa_token_cache():
                     account="account",
                     user="user",
                     password="password",
-                    authenticator="username_password_mfa",
+                    authenticator=authenticator,
                     client_store_temporary_credential=True,
                     client_request_mfa_token=True,
                 ):

test/unit/aio/test_auth_oauth_async.py

@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+import pytest
+
 from snowflake.connector.aio.auth import AuthByOAuth
 
 
@@ -18,6 +20,44 @@ async def test_auth_oauth():
     assert body["data"]["AUTHENTICATOR"] == "OAUTH", body
 
 
+@pytest.mark.parametrize("authenticator", ["oauth", "OAUTH"])
+async def test_oauth_authenticator_is_case_insensitive(monkeypatch, authenticator):
+    """Test that oauth authenticator is case insensitive."""
+    import snowflake.connector.aio
+
+    async def mock_post_request(self, url, headers, json_body, **kwargs):
+        return {
+            "success": True,
+            "message": None,
+            "data": {
+                "token": "TOKEN",
+                "masterToken": "MASTER_TOKEN",
+                "idToken": None,
+                "parameters": [{"name": "SERVICE_NAME", "value": "FAKE_SERVICE_NAME"}],
+            },
+        }
+
+    monkeypatch.setattr(
+        snowflake.connector.aio._network.SnowflakeRestful,
+        "_post_request",
+        mock_post_request,
+    )
+
+    # Create connection with oauth authenticator - OAuth requires a token parameter
+    conn = snowflake.connector.aio.SnowflakeConnection(
+        user="testuser",
+        account="testaccount",
+        authenticator=authenticator,
+        token="test_oauth_token",  # OAuth authentication requires a token
+    )
+    await conn.connect()
+
+    # Verify that the auth_class is an instance of AuthByOAuth
+    assert isinstance(conn.auth_class, AuthByOAuth)
+
+    await conn.close()
+
+
 def test_mro():
     """Ensure that methods from AuthByPluginAsync override those from AuthByPlugin."""
     from snowflake.connector.aio.auth import AuthByPlugin as AuthByPluginAsync