fix test

Author: sfc-gh-yuwang

src/snowflake/snowpark/_internal/data_source/drivers/psycopg2_driver.py

@@ -286,6 +286,7 @@ def process(self, query: str):
                 if session_init_statement is not None:
                     for statement in session_init_statement:
                         cursor.execute(statement)
+                        cursor.fetchall()
                 cursor.execute(query)
                 while True:
                     rows = cursor.fetchmany(fetch_size)

tests/integ/datasource/test_oracledb.py

@@ -260,16 +260,17 @@ def test_query_timeout_and_session_init(session):
         DBMS_LOCK.SLEEP(5);
     END;
 """
-    with pytest.raises(
-        SnowparkDataframeReaderException,
-        match="socket timed out while recovering from previous socket timeout",
-    ):
+    with pytest.raises(SnowparkDataframeReaderException) as error:
         session.read.dbapi(
             create_connection_oracledb,
             table=ORACLEDB_TABLE_NAME,
             query_timeout=1,
             session_init_statement=[statement],
         )
+        assert (
+            "socket timed out while recovering from previous socket timeout" in error
+            or "call timeout of 1000 ms exceeded" in error
+        )
 
 
 def test_query_timeout_and_session_init_udtf(session):

tests/integ/datasource/test_postgres.py

@@ -108,7 +108,7 @@ def test_error_case(session, input_type, input_value, error_message):
 def test_query_timeout_and_session_init(session):
     with pytest.raises(
         SnowparkDataframeReaderException,
-        match=r"due to exception 'QueryCanceled\('canceling statement due to statement timeout",
+        match="canceling statement due to statement timeout",
     ):
         session.read.dbapi(
             create_postgres_connection,

tests/integ/test_data_source_api.py

@@ -1018,7 +1018,7 @@ def fetchmany(self, row_count: int):
         driver.to_snow_type(raw_schema),
         partitions_table,
         "",
-        packages=["pyodbc"],
+        packages=["pyodbc", "snowflake-snowpark-python"],
     )
     Utils.check_answer(df, sql_server_udtf_ingestion_data)