improvements

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()

tests/integration_tests/sql_lab/commands_tests.py

@@ -34,14 +34,15 @@
 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
 
 
 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)