fix: patches Row typing errors when databricks sql imported

apache · Nov 1, 2024 · f372a8c · f372a8c
1 parent f724d25
commit f372a8c
Showing 1 changed file with 23 additions and 20 deletions.
diff --git a/providers/tests/snowflake/operators/test_snowflake_sql.py b/providers/tests/snowflake/operators/test_snowflake_sql.py
@@ -17,6 +17,7 @@
 # under the License.
 from __future__ import annotations
 
+from typing import Any, Callable
 from unittest import mock
 from unittest.mock import MagicMock, patch
 
@@ -27,15 +28,17 @@
 
 databricks = importorskip("databricks")
 
+MockRow = None
 try:
     from databricks.sql.types import Row
 except ImportError:
     # Row is used in the parametrize so it's parsed during collection and we need to have a viable
     # replacement for the collection time when databricks is not installed (Python 3.12 for now)
-    def MockRow(*args, **kwargs):
+    def MockRow(*args: Any, **kwargs: Any) -> MagicMock:
         return MagicMock()
 
-    Row = MockRow
+
+RowType: type[Row] | Callable[..., MagicMock] = Row if "Row" in locals() else MockRow
 
 
 from airflow.models.connection import Connection
@@ -61,59 +64,59 @@ def MockRow(*args, **kwargs):
             "select * from dummy",
             True,
             True,
-            [Row(id=1, value="value1"), Row(id=2, value="value2")],
+            [RowType(id=1, value="value1"), RowType(id=2, value="value2")],
             [[("id",), ("value",)]],
-            ([Row(id=1, value="value1"), Row(id=2, value="value2")]),
+            ([RowType(id=1, value="value1"), RowType(id=2, value="value2")]),
             id="Scalar: Single SQL statement, return_last, split statement",
         ),
         pytest.param(
             "select * from dummy;select * from dummy2",
             True,
             True,
-            [Row(id=1, value="value1"), Row(id=2, value="value2")],
+            [RowType(id=1, value="value1"), RowType(id=2, value="value2")],
             [[("id",), ("value",)]],
-            ([Row(id=1, value="value1"), Row(id=2, value="value2")]),
+            ([RowType(id=1, value="value1"), RowType(id=2, value="value2")]),
             id="Scalar: Multiple SQL statements, return_last, split statement",
         ),
         pytest.param(
             "select * from dummy",
             False,
             False,
-            [Row(id=1, value="value1"), Row(id=2, value="value2")],
+            [RowType(id=1, value="value1"), RowType(id=2, value="value2")],
             [[("id",), ("value",)]],
-            ([Row(id=1, value="value1"), Row(id=2, value="value2")]),
+            ([RowType(id=1, value="value1"), RowType(id=2, value="value2")]),
             id="Scalar: Single SQL statements, no return_last (doesn't matter), no split statement",
         ),
         pytest.param(
             "select * from dummy",
             True,
             False,
-            [Row(id=1, value="value1"), Row(id=2, value="value2")],
+            [RowType(id=1, value="value1"), RowType(id=2, value="value2")],
             [[("id",), ("value",)]],
-            ([Row(id=1, value="value1"), Row(id=2, value="value2")]),
+            ([RowType(id=1, value="value1"), RowType(id=2, value="value2")]),
             id="Scalar: Single SQL statements, return_last (doesn't matter), no split statement",
         ),
         pytest.param(
             ["select * from dummy"],
             False,
             False,
-            [[Row(id=1, value="value1"), Row(id=2, value="value2")]],
+            [[RowType(id=1, value="value1"), RowType(id=2, value="value2")]],
             [[("id",), ("value",)]],
-            [([Row(id=1, value="value1"), Row(id=2, value="value2")])],
+            [([RowType(id=1, value="value1"), RowType(id=2, value="value2")])],
             id="Non-Scalar: Single SQL statements in list, no return_last, no split statement",
         ),
         pytest.param(
             ["select * from dummy", "select * from dummy2"],
             False,
             False,
             [
-                [Row(id=1, value="value1"), Row(id=2, value="value2")],
-                [Row(id2=1, value2="value1"), Row(id2=2, value2="value2")],
+                [RowType(id=1, value="value1"), RowType(id=2, value="value2")],
+                [RowType(id2=1, value2="value1"), RowType(id2=2, value2="value2")],
             ],
             [[("id",), ("value",)], [("id2",), ("value2",)]],
             [
-                ([Row(id=1, value="value1"), Row(id=2, value="value2")]),
-                ([Row(id2=1, value2="value1"), Row(id2=2, value2="value2")]),
+                ([RowType(id=1, value="value1"), RowType(id=2, value="value2")]),
+                ([RowType(id2=1, value2="value1"), RowType(id2=2, value2="value2")]),
             ],
             id="Non-Scalar: Multiple SQL statements in list, no return_last (no matter), no split statement",
         ),
@@ -122,13 +125,13 @@ def MockRow(*args, **kwargs):
             True,
             False,
             [
-                [Row(id=1, value="value1"), Row(id=2, value="value2")],
-                [Row(id2=1, value2="value1"), Row(id2=2, value2="value2")],
+                [RowType(id=1, value="value1"), RowType(id=2, value="value2")],
+                [RowType(id2=1, value2="value1"), RowType(id2=2, value2="value2")],
             ],
             [[("id",), ("value",)], [("id2",), ("value2",)]],
             [
-                ([Row(id=1, value="value1"), Row(id=2, value="value2")]),
-                ([Row(id2=1, value2="value1"), Row(id2=2, value2="value2")]),
+                ([RowType(id=1, value="value1"), RowType(id=2, value="value2")]),
+                ([RowType(id2=1, value2="value1"), RowType(id2=2, value2="value2")]),
             ],
             id="Non-Scalar: Multiple SQL statements in list, return_last (no matter), no split statement",
         ),