feat: refactor all get_sqla_engine to use contextmanager in codebase (

apache#21943)
preset-io · Dec 3, 2022 · 9c06aed · 9c06aed
1 parent fff03b4
commit 9c06aed
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/integration_tests/db_engine_specs/presto_tests.py b/tests/integration_tests/db_engine_specs/presto_tests.py
@@ -82,10 +82,10 @@ def test_get_view_names_without_schema(self):
     def test_get_view_names_without_schema(self):
         database = mock.MagicMock()
         mock_execute = mock.MagicMock()
-        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.execute = (
+        database.get_sqla_engine_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value.execute = (
             mock_execute
         )
-        database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value.fetchall = mock.MagicMock(
+        database.get_sqla_engine_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value.fetchall = mock.MagicMock(
             return_value=[["a", "b,", "c"], ["d", "e"]]
         )
         result = PrestoEngineSpec.get_view_names(database, mock.Mock(), None)