fix snowflake io manager tests on python 3.7 (#7824)

dagster-io · May 10, 2022 · b0e3ca2 · b0e3ca2
1 parent 9ca7fdd
commit b0e3ca2
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/...ter-snowflake-pandas/dagster_snowflake_pandas_tests/test_snowflake_pandas_type_handler.py b/...ter-snowflake-pandas/dagster_snowflake_pandas_tests/test_snowflake_pandas_type_handler.py
@@ -65,8 +65,5 @@ def test_load_input():
                 partition=None,
             ),
         )
-        assert (
-            mock_read_sql.call_args_list[0].kwargs["sql"]
-            == "SELECT * FROM my_db.my_schema.my_table"
-        )
+        assert mock_read_sql.call_args_list[0][1]["sql"] == "SELECT * FROM my_db.my_schema.my_table"
         assert df.equals(DataFrame([{"col1": "a", "col2": 1}]))
diff --git a/python_modules/libraries/dagster-snowflake/dagster_snowflake_tests/test_db_io_manager.py b/python_modules/libraries/dagster-snowflake/dagster_snowflake_tests/test_db_io_manager.py
@@ -53,7 +53,7 @@ def supported_types(self):
 
 def test_asset_out():
     handler = IntHandler()
-    db_client = MagicMock(spec=DbClient)
+    db_client = MagicMock(spec=DbClient, get_select_statement=MagicMock(return_value=""))
     manager = DbIOManager(type_handlers=[handler], db_client=db_client)
     asset_key = AssetKey(["schema1", "table1"])
     output_context = build_output_context(asset_key=asset_key, resource_config=resource_config)
@@ -79,7 +79,7 @@ def test_asset_out():
 
 def test_asset_out_columns():
     handler = IntHandler()
-    db_client = MagicMock(spec=DbClient)
+    db_client = MagicMock(spec=DbClient, get_select_statement=MagicMock(return_value=""))
     manager = DbIOManager(type_handlers=[handler], db_client=db_client)
     asset_key = AssetKey(["schema1", "table1"])
     output_context = build_output_context(asset_key=asset_key, resource_config=resource_config)
@@ -107,7 +107,7 @@ def test_asset_out_columns():
 
 def test_asset_out_partitioned():
     handler = IntHandler()
-    db_client = MagicMock(spec=DbClient)
+    db_client = MagicMock(spec=DbClient, get_select_statement=MagicMock(return_value=""))
     manager = DbIOManager(type_handlers=[handler], db_client=db_client)
     asset_key = AssetKey(["schema1", "table1"])
     output_context = MagicMock(
@@ -148,7 +148,7 @@ def test_asset_out_partitioned():
 def test_different_output_and_input_types():
     int_handler = IntHandler()
     str_handler = StringHandler()
-    db_client = MagicMock(spec=DbClient)
+    db_client = MagicMock(spec=DbClient, get_select_statement=MagicMock(return_value=""))
     manager = DbIOManager(type_handlers=[int_handler, str_handler], db_client=db_client)
     asset_key = AssetKey(["schema1", "table1"])
     output_context = build_output_context(asset_key=asset_key, resource_config=resource_config)
@@ -176,7 +176,7 @@ def test_different_output_and_input_types():
 
 def test_non_asset_out():
     handler = IntHandler()
-    db_client = MagicMock(spec=DbClient)
+    db_client = MagicMock(spec=DbClient, get_select_statement=MagicMock(return_value=""))
     manager = DbIOManager(type_handlers=[handler], db_client=db_client)
     output_context = build_output_context(
         name="table1", metadata={"schema": "schema1"}, resource_config=resource_config