[async] adjust async wif tests (#2510)

Author: sfc-gh-pczajka

ci/wif/test_wif.sh

@@ -5,6 +5,13 @@ set -o pipefail
 export SF_OCSP_TEST_MODE=true
 export RUN_WIF_TESTS=true
 
+#setup pytest
+/opt/python/cp39-cp39/bin/python -m pip install --break-system-packages pytest pytest-asyncio
+
+# test WIF without asyncio installed
+/opt/python/cp39-cp39/bin/python -m pip install --break-system-packages -e .
+/opt/python/cp39-cp39/bin/python -m pytest test/wif/* --ignore test/wif/*_async.py
+
+# test WIF with asyncio installed
 /opt/python/cp39-cp39/bin/python -m pip install --break-system-packages -e '.[aio]'
-/opt/python/cp39-cp39/bin/python -m pip install --break-system-packages pytest
 /opt/python/cp39-cp39/bin/python -m pytest test/wif/*

test/wif/test_wif_async.py

@@ -22,6 +22,9 @@
 ACCOUNT = os.getenv("SNOWFLAKE_TEST_WIF_ACCOUNT")
 HOST = os.getenv("SNOWFLAKE_TEST_WIF_HOST")
 PROVIDER = os.getenv("SNOWFLAKE_TEST_WIF_PROVIDER")
+EXPECTED_USERNAME = os.getenv("SNOWFLAKE_TEST_WIF_USERNAME")
+IMPERSONATION_PATH = os.getenv("SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH")
+EXPECTED_USERNAME_IMPERSONATION = os.getenv("SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION")
 
 
 @pytest.mark.wif
@@ -34,8 +37,8 @@ async def test_wif_defined_provider_async():
         "workload_identity_provider": PROVIDER,
     }
     assert await connect_and_execute_simple_query_async(
-        connection_params
-    ), "Failed to connect with using WIF - automatic provider detection"
+        connection_params, EXPECTED_USERNAME
+    ), f"Failed to connect with using WIF using provider {PROVIDER}"
 
 
 @pytest.mark.wif
@@ -53,17 +56,47 @@ async def test_should_authenticate_using_oidc_async():
     }
 
     assert await connect_and_execute_simple_query_async(
-        connection_params
+        connection_params, expected_user=None
     ), "Failed to connect using WIF with OIDC provider"
 
 
-async def connect_and_execute_simple_query_async(connection_params) -> bool:
+@pytest.mark.wif
+@pytest.mark.aio
+@pytest.mark.skip("Impersonation is still being developed")
+async def test_should_authenticate_with_impersonation_async():
+    if not isinstance(IMPERSONATION_PATH, str) or not IMPERSONATION_PATH:
+        pytest.skip("Skipping test - IMPERSONATION_PATH is not set")
+
+    logger.debug(f"Using impersonation path: {IMPERSONATION_PATH}")
+    impersonation_path_list = IMPERSONATION_PATH.split(",")
+
+    connection_params = {
+        "host": HOST,
+        "account": ACCOUNT,
+        "authenticator": "WORKLOAD_IDENTITY",
+        "workload_identity_provider": PROVIDER,
+        "workload_identity_impersonation_path": impersonation_path_list,
+    }
+
+    assert await connect_and_execute_simple_query_async(
+        connection_params, EXPECTED_USERNAME_IMPERSONATION
+    ), f"Failed to connect using WIF with provider {PROVIDER}"
+
+
+async def connect_and_execute_simple_query_async(
+    connection_params, expected_user=None
+) -> bool:
     try:
         logger.info("Trying to connect to Snowflake")
         async with snowflake.connector.aio.connect(**connection_params) as con:
-            result = await con.cursor().execute("select 1;")
-            logger.debug(await result.fetchall())
-            logger.info("Successfully connected to Snowflake")
+            result = await con.cursor().execute("select current_user();")
+            (user,) = await result.fetchone()
+            logger.debug(user)
+            if expected_user:
+                assert (
+                    expected_user == user
+                ), f"Expected user '{expected_user}', got user '{user}'"
+            logger.info(f"Successfully connected to Snowflake as {user}")
             return True
     except Exception as e:
         logger.error(e)