SNOW-690343: Fixing import errors for olddriver tests (#1321)

sfc-gh-heshah · web-flow · commit d35ffe8ed331 · 2022-11-07T11:53:10.000-08:00
diff --git a/test/integ/pandas/test_pandas_tools.py b/test/integ/pandas/test_pandas_tools.py
@@ -14,7 +14,11 @@
 
 from snowflake.connector import DictCursor
 from snowflake.connector.errors import ProgrammingError
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ...randomize import random_string
 
 from ...lazy_var import LazyVar
 
diff --git a/test/integ/test_bindings.py b/test/integ/test_bindings.py
@@ -20,7 +20,11 @@
 
 from snowflake.connector.converter import convert_datetime_to_epoch
 from snowflake.connector.errors import ForbiddenError, ProgrammingError
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 tempfile.gettempdir()
 
diff --git a/test/integ/test_cursor.py b/test/integ/test_cursor.py
@@ -51,7 +51,11 @@ class ResultMetadata(NamedTuple):
 )
 from snowflake.connector.sqlstate import SQLSTATE_FEATURE_NOT_SUPPORTED
 from snowflake.connector.telemetry import TelemetryField
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 try:
     from snowflake.connector.constants import (
diff --git a/test/integ/test_dataintegrity.py b/test/integ/test_dataintegrity.py
@@ -18,7 +18,11 @@
 import pytz
 
 from snowflake.connector.dbapi import DateFromTicks, TimeFromTicks, TimestampFromTicks
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 
 def table_exists(conn_cnx, name):
diff --git a/test/integ/test_dbapi.py b/test/integ/test_dbapi.py
@@ -17,7 +17,11 @@
 import snowflake.connector
 import snowflake.connector.dbapi
 from snowflake.connector import dbapi, errorcode, errors
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 TABLE1 = "dbapi_ddl1"
 TABLE2 = "dbapi_ddl2"
diff --git a/test/integ/test_put_get.py b/test/integ/test_put_get.py
@@ -17,7 +17,11 @@
 import pytest
 
 from snowflake.connector import OperationalError
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 from ..generate_test_files import generate_k_lines_of_n_files
 from ..integ_helpers import put
diff --git a/test/integ/test_put_get_compress_enc.py b/test/integ/test_put_get_compress_enc.py
@@ -11,7 +11,10 @@
 
 import pytest
 
-from snowflake.connector.util_text import random_string
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 from ..integ_helpers import put
 
diff --git a/test/integ/test_put_get_medium.py b/test/integ/test_put_get_medium.py
@@ -23,7 +23,11 @@
     SnowflakeProgressPercentage,
     SnowflakeS3ProgressPercentage,
 )
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 from ..generate_test_files import generate_k_lines_of_n_files
 from ..integ_helpers import put
diff --git a/test/integ/test_put_get_user_stage.py b/test/integ/test_put_get_user_stage.py
@@ -15,7 +15,11 @@
 import pytest
 
 from snowflake.connector.cursor import SnowflakeCursor
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 from ..generate_test_files import generate_k_lines_of_n_files
 from ..integ_helpers import put
diff --git a/test/integ/test_put_get_with_aws_token.py b/test/integ/test_put_get_with_aws_token.py
@@ -28,7 +28,10 @@
 except ImportError:
     pass
 
-from snowflake.connector.util_text import random_string
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 from ..integ_helpers import put
 
diff --git a/test/integ/test_put_get_with_azure_token.py b/test/integ/test_put_get_with_azure_token.py
@@ -19,7 +19,11 @@
     SnowflakeAzureProgressPercentage,
     SnowflakeProgressPercentage,
 )
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 from ..generate_test_files import generate_k_lines_of_n_files
 from ..integ_helpers import put
diff --git a/test/integ/test_put_get_with_gcp_account.py b/test/integ/test_put_get_with_gcp_account.py
@@ -33,7 +33,10 @@
 
     SnowflakeGCSRestClient = None
 
-from snowflake.connector.util_text import random_string
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 from ..generate_test_files import generate_k_lines_of_n_files
 from ..integ_helpers import put
diff --git a/test/randomize.py b/test/randomize.py
@@ -0,0 +1,30 @@
+#
+# Copyright (c) 2012-2021 Snowflake Computing Inc. All rights reserved.
+#
+"""
+This module was added back to the repository for compatibility with the old driver tests that
+rely on random_string from this file for functionality.
+"""
+
+from __future__ import annotations
+
+import random
+import string
+from typing import Sequence
+
+
+def random_string(
+    length: int,
+    prefix: str = "",
+    suffix: str = "",
+    choices: Sequence[str] = string.ascii_lowercase,
+) -> str:
+    """Our convenience function to generate random string for object names.
+    Args:
+        length: How many random characters to choose from choices.
+        prefix: Prefix to add to random string generated.
+        suffix: Suffix to add to random string generated.
+        choices: A generator of things to choose from.
+    """
+    random_part = "".join([random.choice(choices) for _ in range(length)])
+    return "".join([prefix, random_part, suffix])
diff --git a/test/unit/test_gcs_client.py b/test/unit/test_gcs_client.py
@@ -14,7 +14,11 @@
 
 from snowflake.connector import SnowflakeConnection
 from snowflake.connector.constants import SHA256_DIGEST, ResultStatus
-from snowflake.connector.util_text import random_string
+
+try:
+    from snowflake.connector.util_text import random_string
+except ImportError:
+    from ..randomize import random_string
 
 try:
     from snowflake.connector.errors import RequestExceedMaxRetryError
diff --git a/test/unit/test_ocsp.py b/test/unit/test_ocsp.py
@@ -18,12 +18,7 @@
 from snowflake.connector import OperationalError
 from snowflake.connector.errors import RevocationCheckError
 from snowflake.connector.ocsp_asn1crypto import SnowflakeOCSPAsn1Crypto as SFOCSP
-from snowflake.connector.ocsp_snowflake import (
-    OCSPCache,
-    OCSPResponseValidationResult,
-    OCSPServer,
-    SnowflakeOCSP,
-)
+from snowflake.connector.ocsp_snowflake import OCSPCache, OCSPServer, SnowflakeOCSP
 from snowflake.connector.ssl_wrap_socket import _openssl_connect
 
 try:
@@ -213,7 +208,10 @@ def test_ocsp_with_file_cache(tmpdir):
         assert ocsp.validate(url, connection), f"Failed to validate: {url}"
 
 
+@pytest.mark.skipolddriver
 def test_ocsp_with_bogus_cache_files(tmpdir):
+    from snowflake.connector.ocsp_snowflake import OCSPResponseValidationResult
+
     """Attempts to use bogus OCSP response data."""
     cache_file_name, target_hosts = _store_cache_in_file(tmpdir)
 
@@ -245,7 +243,10 @@ def test_ocsp_with_bogus_cache_files(tmpdir):
         )
 
 
+@pytest.mark.skipolddriver
 def test_ocsp_with_outdated_cache(tmpdir):
+    from snowflake.connector.ocsp_snowflake import OCSPResponseValidationResult
+
     """Attempts to use outdated OCSP response cache file."""
     cache_file_name, target_hosts = _store_cache_in_file(tmpdir)
 
