SNOW-2235955: adding MFA test in Python

ci/test_authentication.sh

@@ -23,5 +23,5 @@ docker run \
   -v $(cd $THIS_DIR/.. && pwd):/mnt/host \
   -v $WORKSPACE:/mnt/workspace \
   --rm \
-    nexus.int.snowflakecomputing.com:8086/docker/snowdrivers-test-external-browser-python:2 \
+    nexus.int.snowflakecomputing.com:8086/docker/snowdrivers-test-external-browser-python:3 \
   "/mnt/host/ci/container/test_authentication.sh"

test/auth/authorization_parameters.py

@@ -77,6 +77,13 @@ def __init__(self):
     def get_base_connection_parameters(self) -> dict[str, Union[str, bool, int]]:
         return self.basic_config
 
+    def get_mfa_connection_parameters(self) -> dict[str, Union[str, bool, int]]:
+        config = self.basic_config.copy()
+        config["user"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_MFA_USER")
+        config["password"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_MFA_PASSWORD")
+        config["authenticator"] = "USERNAME_PASSWORD_MFA"
+        return config
+
     def get_key_pair_connection_parameters(self):
         config = self.basic_config.copy()
         config["authenticator"] = "KEY_PAIR_AUTHENTICATOR"

test/auth/authorization_test_helper.py

@@ -154,6 +154,25 @@ def _provide_credentials(self, scenario: Scenario, login: str, password: str):
             self.error_msg = e
             raise RuntimeError(e)
 
+    def get_totp(self, seed: str = "") -> []:
+        if self.auth_test_env == "docker":
+            try:
+                provide_totp_generator_path = "/externalbrowser/totpGenerator.js"
+                process = subprocess.run(
+                    ["node", provide_totp_generator_path, seed],
+                    timeout=40,
+                    capture_output=True,
+                    text=True,
+                )
+                logger.debug(f"OUTPUT:  {process.stdout}, ERRORS: {process.stderr}")
+                return process.stdout.strip().split()
+            except Exception as e:
+                self.error_msg = e
+                raise RuntimeError(e)
+        else:
+            logger.info("TOTP generation is not supported in this environment")
+            return ""
+
     def connect_using_okta_connection_and_execute_custom_command(
         self, command: str, return_token: bool = False
     ) -> Union[bool, str]:

test/auth/test_mfa.py

@@ -0,0 +1,75 @@
+import logging
+from test.auth.authorization_parameters import AuthConnectionParameters
+from test.auth.authorization_test_helper import AuthorizationTestHelper
+
+import pytest
+
+
+@pytest.mark.auth
+def test_mfa_successful():
+    connection_parameters = AuthConnectionParameters().get_mfa_connection_parameters()
+    connection_parameters["client_request_mfa_token"] = True
+    test_helper = AuthorizationTestHelper(connection_parameters)
+    totp_codes = test_helper.get_totp()
+    logging.info(f"Got {len(totp_codes)} TOTP codes to try")
+
+    connection_success = False
+    last_error = ""
+
+    # Try each TOTP code until one works
+    for i, totp_code in enumerate(totp_codes):
+        logging.info(f"Trying TOTP code {i+1}/{len(totp_codes)}")
+
+        # Update the passcode in connection parameters
+        connection_parameters["passcode"] = totp_code
+        test_helper.update_config(connection_parameters)
+
+        # Clear any previous error
+        test_helper.error_msg = ""
+
+        # Try to connect
+        connection_success = test_helper.connect_and_execute_simple_query()
+
+        if connection_success:
+            logging.info(f"Successfully connected with TOTP code {i+1}")
+            break
+        else:
+            last_error = str(test_helper.error_msg)
+            logging.warning(f"TOTP code {i+1} failed: {last_error}")
+
+            # Check if it's a TOTP-related error before continuing
+            if "TOTP Invalid" in last_error:
+                logging.info("TOTP/MFA error detected, trying next code...")
+                continue
+            else:
+                # If it's not a TOTP error, no point trying other codes
+                logging.error(f"Non-TOTP error detected: {last_error}")
+                break
+
+    assert (
+        connection_success
+    ), f"Failed to connect with any of the {len(totp_codes)} TOTP codes. Last error: {last_error}"
+    assert (
+        test_helper.error_msg == ""
+    ), f"Final error message should be empty but got: {test_helper.error_msg}"
+
+    logging.info("Testing MFA token caching with second connection...")
+
+    # Create fresh connection parameters for cache test - but use same exact config
+    cache_test_parameters = connection_parameters.copy()
+
+    # Remove the passcode that was set in the loop - should use cached MFA token instead
+    if "passcode" in cache_test_parameters:
+        del cache_test_parameters["passcode"]
+
+    # Create new helper for cache test
+    cache_test_helper = AuthorizationTestHelper(cache_test_parameters)
+    cache_test_helper.error_msg = ""
+    cache_connection_success = cache_test_helper.connect_and_execute_simple_query()
+
+    assert (
+        cache_connection_success
+    ), f"Failed to connect with cached MFA token. Error: {cache_test_helper.error_msg}"
+    assert (
+        cache_test_helper.error_msg == ""
+    ), f"Cache test error message should be empty but got: {cache_test_helper.error_msg}"