SNOW-199087 Moved functions using pytest to test folder (#1362)

Co-authored-by: Mark Keller <[email protected]>

Author: sfc-gh-heshah

src/snowflake/connector/test_util.py

@@ -7,12 +7,6 @@
 
 import logging
 import os
-import time
-
-import pytest
-
-import snowflake.connector.connection
-from snowflake.connector.constants import QueryStatus
 
 from .compat import IS_LINUX
 
@@ -35,37 +29,3 @@
         )
     )
     rt_plain_logger.addHandler(ch)
-
-
-def _wait_while_query_running(
-    con: snowflake.connector.connection.SnowflakeConnection,
-    sfqid: str,
-    sleep_time: int,
-    dont_cache: bool = False,
-) -> None:
-    """
-    Checks if the provided still returns that it is still running, and if so,
-    sleeps for the specified time in a while loop.
-    """
-    query_status = con._get_query_status if dont_cache else con.get_query_status
-    while con.is_still_running(query_status(sfqid)):
-        time.sleep(sleep_time)
-
-
-def _wait_until_query_success(
-    con: snowflake.connector.connection.SnowflakeConnection,
-    sfqid: str,
-    num_checks: int,
-    sleep_per_check: int,
-) -> None:
-    for _ in range(num_checks):
-        status = con.get_query_status(sfqid)
-        if status == QueryStatus.SUCCESS:
-            break
-        time.sleep(sleep_per_check)
-    else:
-        pytest.fail(
-            "We should have broke out of wait loop for query success."
-            f"Query ID: {sfqid}"
-            f"Final query status: {status}"
-        )

test/helpers.py

@@ -4,11 +4,22 @@
 #
 from __future__ import annotations
 
-from typing import Pattern, Sequence
+import time
+from typing import TYPE_CHECKING, Pattern, Sequence
 from unittest.mock import Mock
 
+import pytest
+
 from snowflake.connector.compat import OK
 
+if TYPE_CHECKING:
+    import snowflake.connector.connection
+
+try:
+    from snowflake.connector.constants import QueryStatus
+except ImportError:
+    QueryStatus = None
+
 
 def create_mock_response(status_code: int) -> Mock:
     """Create a Mock "Response" with a given status code. See `test_result_batch.py` for examples.
@@ -43,3 +54,37 @@ def verify_log_tuple(
             ):
                 return True
     return False
+
+
+def _wait_while_query_running(
+    con: snowflake.connector.connection.SnowflakeConnection,
+    sfqid: str,
+    sleep_time: int,
+    dont_cache: bool = False,
+) -> None:
+    """
+    Checks if the provided still returns that it is still running, and if so,
+    sleeps for the specified time in a while loop.
+    """
+    query_status = con._get_query_status if dont_cache else con.get_query_status
+    while con.is_still_running(query_status(sfqid)):
+        time.sleep(sleep_time)
+
+
+def _wait_until_query_success(
+    con: snowflake.connector.connection.SnowflakeConnection,
+    sfqid: str,
+    num_checks: int,
+    sleep_per_check: int,
+) -> None:
+    for _ in range(num_checks):
+        status = con.get_query_status(sfqid)
+        if status == QueryStatus.SUCCESS:
+            break
+        time.sleep(sleep_per_check)
+    else:
+        pytest.fail(
+            "We should have broke out of wait loop for query success."
+            f"Query ID: {sfqid}"
+            f"Final query status: {status}"
+        )

test/integ/test_multi_statement.py

@@ -6,6 +6,8 @@
 
 from snowflake.connector.version import VERSION
 
+from ..helpers import _wait_until_query_success, _wait_while_query_running
+
 pytestmark = [
     pytest.mark.skipolddriver,
     pytest.mark.xfail(
@@ -17,14 +19,6 @@
 import snowflake.connector.cursor
 from snowflake.connector import ProgrammingError, errors
 
-try:
-    from snowflake.connector.test_util import (
-        _wait_until_query_success,
-        _wait_while_query_running,
-    )
-except ModuleNotFoundError:
-    pass
-
 try:  # pragma: no cover
     from snowflake.connector.constants import (
         PARAMETER_MULTI_STATEMENT_COUNT,