From d48e9e703cf6e1b31af1c6420cf6dd941648cb2c Mon Sep 17 00:00:00 2001 From: Diego Medina Date: Wed, 8 Feb 2023 23:09:50 -0300 Subject: [PATCH] improvements --- tests/integration_tests/sql_lab/api_tests.py | 9 ++++----- tests/integration_tests/sql_lab/commands_tests.py | 9 +++------ 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/tests/integration_tests/sql_lab/api_tests.py b/tests/integration_tests/sql_lab/api_tests.py index 23b6f2c9deec0..93beb380f0db6 100644 --- a/tests/integration_tests/sql_lab/api_tests.py +++ b/tests/integration_tests/sql_lab/api_tests.py @@ -185,9 +185,7 @@ def test_get_results_with_display_limit(self): def test_export_results(self, get_df_mock: mock.Mock) -> None: self.login() - database = Database( - database_name="my_export_database", sqlalchemy_uri="sqlite://" - ) + database = get_example_database() query_obj = Query( client_id="test", database=database, @@ -203,8 +201,8 @@ def test_export_results(self, get_df_mock: mock.Mock) -> None: results_key="test_abc", ) - db.session.add(database) db.session.add(query_obj) + db.session.commit() get_df_mock.return_value = pd.DataFrame({"foo": [1, 2, 3]}) @@ -213,4 +211,5 @@ def test_export_results(self, get_df_mock: mock.Mock) -> None: expected_data = csv.reader(io.StringIO("foo\n1\n2")) self.assertEqual(list(expected_data), list(data)) - db.session.rollback() + db.session.delete(query_obj) + db.session.commit() diff --git a/tests/integration_tests/sql_lab/commands_tests.py b/tests/integration_tests/sql_lab/commands_tests.py index 2293b20a5fcbb..edb71552370b7 100644 --- a/tests/integration_tests/sql_lab/commands_tests.py +++ b/tests/integration_tests/sql_lab/commands_tests.py @@ -34,6 +34,7 @@ from superset.sqllab.commands import export, results from superset.sqllab.limiting_factor import LimitingFactor from superset.utils import core as utils +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase @@ -41,7 +42,7 @@ class TestSqlResultExportCommand(SupersetTestCase): @pytest.fixture() def create_database_and_query(self): with self.create_app().app_context(): - database = Database(database_name="my_database", sqlalchemy_uri="sqlite://") + database = get_example_database() query_obj = Query( client_id="test", database=database, @@ -57,14 +58,12 @@ def create_database_and_query(self): results_key="abc_query", ) - db.session.add(database) db.session.add(query_obj) db.session.commit() yield db.session.delete(query_obj) - db.session.delete(database) db.session.commit() @pytest.mark.usefixtures("create_database_and_query") @@ -193,7 +192,7 @@ class TestSqlExecutionResultsCommand(SupersetTestCase): @pytest.fixture() def create_database_and_query(self): with self.create_app().app_context(): - database = Database(database_name="my_database", sqlalchemy_uri="sqlite://") + database = get_example_database() query_obj = Query( client_id="test", database=database, @@ -209,14 +208,12 @@ def create_database_and_query(self): results_key="abc_query", ) - db.session.add(database) db.session.add(query_obj) db.session.commit() yield db.session.delete(query_obj) - db.session.delete(database) db.session.commit() @patch("superset.sqllab.commands.results.results_backend_use_msgpack", False)