QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-19 16:14:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore(ssh-tunnel): Refactor establishing raw connection with contextmanger (#22366)

Browse Source
This commit is contained in:
Hugh A. Miles II
2022-12-29 14:20:45 -05:00
committed by GitHub
parent b6d39d194c
commit 9c0d6c51f1
9 changed files with 112 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
tests/integration_tests/db_engine_specs/presto_tests.py
View File

@@ -37,10 +37,8 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
def test_get_view_names_with_schema(self):
database = mock.MagicMock()
mock_execute = mock.MagicMock()
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_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value.fetchall = mock.MagicMock(
database.get_raw_connection().__enter__().cursor().execute = mock_execute
database.get_raw_connection().__enter__().cursor().fetchall = mock.MagicMock(
return_value=[["a", "b,", "c"], ["d", "e"]]
)
@@ -61,10 +59,8 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
def test_get_view_names_without_schema(self):
database = mock.MagicMock()
mock_execute = mock.MagicMock()
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_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value.fetchall = mock.MagicMock(
database.get_raw_connection().__enter__().cursor().execute = mock_execute
database.get_raw_connection().__enter__().cursor().fetchall = mock.MagicMock(
return_value=[["a", "b,", "c"], ["d", "e"]]
)
result = PrestoEngineSpec.get_view_names(database, mock.Mock(), None)
@@ -823,15 +819,9 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
mock_execute = mock.MagicMock()
mock_fetchall = mock.MagicMock(return_value=[["a", "b,", "c"], ["d", "e"]])
database = mock.MagicMock()
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_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value.fetchall = (
mock_fetchall
)
database.get_sqla_engine_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value.poll.return_value = (
False
)
database.get_raw_connection().__enter__().cursor().execute = mock_execute
database.get_raw_connection().__enter__().cursor().fetchall = mock_fetchall
database.get_raw_connection().__enter__().cursor().return_value = False
schema = "schema"
table = "table"
result = PrestoEngineSpec.get_create_view(database, schema=schema, table=table)
@@ -841,9 +831,7 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
def test_get_create_view_exception(self):
mock_execute = mock.MagicMock(side_effect=Exception())
database = mock.MagicMock()
database.get_sqla_engine_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value.execute = (
mock_execute
)
database.get_raw_connection().__enter__().cursor().execute = mock_execute
schema = "schema"
table = "table"
with self.assertRaises(Exception):
@@ -854,9 +842,7 @@ class TestPrestoDbEngineSpec(TestDbEngineSpec):
mock_execute = mock.MagicMock(side_effect=DatabaseError())
database = mock.MagicMock()
database.get_sqla_engine_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value.execute = (
mock_execute
)
database.get_raw_connection().__enter__().cursor().execute = mock_execute
schema = "schema"
table = "table"
result = PrestoEngineSpec.get_create_view(database, schema=schema, table=table)