NO-SNOW: fix dbpi tests (#3399)

Author: sfc-gh-aling

.github/workflows/daily_precommit.yml

@@ -154,6 +154,8 @@ jobs:
           TOX_PARALLEL_NO_SPINNER: 1
         shell: bash
       - name: Run data source tests
+        # psycopg2 is not supported on macos 3.9
+        if: ${{ !(matrix.os == 'macos-latest' && matrix.python-version == '3.9') }}
         run: python -m tox -e datasource
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}

tests/integ/datasource/test_databricks.py

@@ -33,7 +33,7 @@
     TEST_TABLE_NAME,
     DATABRICKS_TEST_EXTERNAL_ACCESS_INTEGRATION,
 )
-from tests.utils import IS_IN_STORED_PROC
+from tests.utils import IS_IN_STORED_PROC, Utils
 
 DEPENDENCIES_PACKAGE_UNAVAILABLE = True
 try:
@@ -65,15 +65,16 @@ def test_basic_databricks(session, input_type, input_value):
         input_type: input_value,
     }
     df = session.read.dbapi(create_databricks_connection, **input_dict).order_by(
-        "COL_BYTE"
+        "COL_BYTE", ascending=True
     )
-    ret = df.collect()
-    assert ret == EXPECTED_TEST_DATA and df.schema == EXPECTED_TYPE
+    Utils.check_answer(df, EXPECTED_TEST_DATA)
+    assert df.schema == EXPECTED_TYPE
 
     table_name = random_name_for_temp_object(TempObjectType.TABLE)
     df.write.save_as_table(table_name, mode="overwrite", table_type="temp")
-    df2 = session.table(table_name).order_by("COL_BYTE")
-    assert df2.collect() == EXPECTED_TEST_DATA and df2.schema == EXPECTED_TYPE
+    df2 = session.table(table_name).order_by("COL_BYTE", ascending=True)
+    Utils.check_answer(df2, EXPECTED_TEST_DATA)
+    assert df2.schema == EXPECTED_TYPE
 
 
 @pytest.mark.parametrize(
@@ -146,9 +147,9 @@ def local_create_databricks_connection():
         udtf_configs={
             "external_access_integration": DATABRICKS_TEST_EXTERNAL_ACCESS_INTEGRATION
         },
-    ).order_by("COL_BYTE")
-    ret = df.collect()
-    assert ret == EXPECTED_TEST_DATA and df.schema == EXPECTED_TYPE
+    ).order_by("COL_BYTE", ascending=True)
+    Utils.check_answer(df, EXPECTED_TEST_DATA)
+    assert df.schema == EXPECTED_TYPE
 
     assert (
         "TEMPORARY  FUNCTION  data_source_udtf_" "" in caplog.text

tests/resources/test_data_source_dir/test_databricks_data.py

@@ -22,7 +22,7 @@
     VariantType,
 )
 
-TEST_TABLE_NAME = "ALL_TYPE_TABLE_2"  # ALL_TYPE_TABLE_2 contains None data while ALL_TYPE_TABLE doesn't
+TEST_TABLE_NAME = "ALL_TYPE_TABLE_3"  # ALL_TYPE_TABLE_3 contains None data while ALL_TYPE_TABLE doesn't
 TZ_INFO = pytz.timezone("America/Los_Angeles")
 EXPECTED_TEST_DATA = [
     tuple([None] * 18),
@@ -44,7 +44,7 @@
         datetime.datetime(2025, 4, 16, 17, 51, 18, 565000),
         "[\n  6,\n  87\n]",
         '{\n  "key1": 83,\n  "key2": 12\n}',
-        '{\n  "field1": "f_77",\n  "field2": 13\n}',
+        '{\n  "field1": "f_77",\n  "field2": "13"\n}',
         "1-8",
         "0 23:14:09.000000000",
     ),
@@ -66,7 +66,7 @@
         datetime.datetime(2025, 4, 16, 17, 42, 39, 565000),
         "[\n  0,\n  89\n]",
         '{\n  "key1": 97,\n  "key2": 33\n}',
-        '{\n  "field1": "f_84",\n  "field2": 1\n}',
+        '{\n  "field1": "f_84",\n  "field2": "1"\n}',
         "1-10",
         "2 11:12:05.000000000",
     ),
@@ -88,7 +88,7 @@
         datetime.datetime(2025, 4, 16, 17, 49, 8, 565000),
         "[\n  82,\n  40\n]",
         '{\n  "key1": 71,\n  "key2": 81\n}',
-        '{\n  "field1": "f_25",\n  "field2": 25\n}',
+        '{\n  "field1": "f_25",\n  "field2": "25"\n}',
         "3-10",
         "18 14:29:08.000000000",
     ),
@@ -110,7 +110,7 @@
         datetime.datetime(2025, 4, 16, 17, 47, 9, 565000),
         "[\n  81,\n  65\n]",
         '{\n  "key1": 67,\n  "key2": 88\n}',
-        '{\n  "field1": "f_98",\n  "field2": 69\n}',
+        '{\n  "field1": "f_98",\n  "field2": "69"\n}',
         "0-7",
         "19 06:25:08.000000000",
     ),
@@ -132,7 +132,7 @@
         datetime.datetime(2025, 4, 16, 17, 50, 8, 565000),
         "[\n  92,\n  27\n]",
         '{\n  "key1": 52,\n  "key2": 65\n}',
-        '{\n  "field1": "f_85",\n  "field2": 50\n}',
+        '{\n  "field1": "f_85",\n  "field2": "50"\n}',
         "7-4",
         "22 04:52:41.000000000",
     ),