From 924f0520351a6a33ac16bc372555ecf2903caa28 Mon Sep 17 00:00:00 2001 From: DShi Date: Wed, 30 Oct 2024 11:26:49 -0400 Subject: [PATCH] fix: patches Row typing errors when databricks sql imported --- .../snowflake/operators/test_snowflake_sql.py | 43 ++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/providers/tests/snowflake/operators/test_snowflake_sql.py b/providers/tests/snowflake/operators/test_snowflake_sql.py index 47dcf27122df1..e6c6dac281655 100644 --- 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,45 +64,45 @@ 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( @@ -107,13 +110,13 @@ def MockRow(*args, **kwargs): 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", ),