Revert "chore: Cleanup database sessions (#10427)" (#10537)

This reverts commit 7645fc85c3.

tests/alerts_tests.py

@@ -32,118 +32,112 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def setup_module():
+@pytest.yield_fixture(scope="module")
+def setup_database():
     with app.app_context():
         slice_id = db.session.query(Slice).all()[0].id
         database_id = utils.get_example_database().id
 
-        alerts = [
-            Alert(
-                id=1,
-                label="alert_1",
-                active=True,
-                crontab="*/1 * * * *",
-                sql="SELECT 0",
-                alert_type="email",
-                slice_id=slice_id,
-                database_id=database_id,
-            ),
-            Alert(
-                id=2,
-                label="alert_2",
-                active=True,
-                crontab="*/1 * * * *",
-                sql="SELECT 55",
-                alert_type="email",
-                slice_id=slice_id,
-                database_id=database_id,
-            ),
-            Alert(
-                id=3,
-                label="alert_3",
-                active=False,
-                crontab="*/1 * * * *",
-                sql="UPDATE 55",
-                alert_type="email",
-                slice_id=slice_id,
-                database_id=database_id,
-            ),
-            Alert(id=4, active=False, label="alert_4", database_id=-1),
-            Alert(id=5, active=False, label="alert_5", database_id=database_id),
-        ]
+        alert1 = Alert(
+            id=1,
+            label="alert_1",
+            active=True,
+            crontab="*/1 * * * *",
+            sql="SELECT 0",
+            alert_type="email",
+            slice_id=slice_id,
+            database_id=database_id,
+        )
+        alert2 = Alert(
+            id=2,
+            label="alert_2",
+            active=True,
+            crontab="*/1 * * * *",
+            sql="SELECT 55",
+            alert_type="email",
+            slice_id=slice_id,
+            database_id=database_id,
+        )
+        alert3 = Alert(
+            id=3,
+            label="alert_3",
+            active=False,
+            crontab="*/1 * * * *",
+            sql="UPDATE 55",
+            alert_type="email",
+            slice_id=slice_id,
+            database_id=database_id,
+        )
+        alert4 = Alert(id=4, active=False, label="alert_4", database_id=-1)
+        alert5 = Alert(id=5, active=False, label="alert_5", database_id=database_id)
 
-        db.session.bulk_save_objects(alerts)
+        for num in range(1, 6):
+            eval(f"db.session.add(alert{num})")
         db.session.commit()
+        yield db.session
 
-
-def teardown_module():
-    with app.app_context():
         db.session.query(AlertLog).delete()
         db.session.query(Alert).delete()
 
 
 @patch("superset.tasks.schedules.deliver_alert")
 @patch("superset.tasks.schedules.logging.Logger.error")
-def test_run_alert_query(mock_error, mock_deliver_alert):
-    with app.app_context():
-        run_alert_query(db.session.query(Alert).filter_by(id=1).one().id)
-        alert1 = db.session.query(Alert).filter_by(id=1).one()
-        assert mock_deliver_alert.call_count == 0
-        assert len(alert1.logs) == 1
-        assert alert1.logs[0].alert_id == 1
-        assert alert1.logs[0].state == "pass"
+def test_run_alert_query(mock_error, mock_deliver, setup_database):
+    database = setup_database
+    run_alert_query(database.query(Alert).filter_by(id=1).one().id, database)
+    alert1 = database.query(Alert).filter_by(id=1).one()
+    assert mock_deliver.call_count == 0
+    assert len(alert1.logs) == 1
+    assert alert1.logs[0].alert_id == 1
+    assert alert1.logs[0].state == "pass"
 
-        run_alert_query(db.session.query(Alert).filter_by(id=2).one().id)
-        alert2 = db.session.query(Alert).filter_by(id=2).one()
-        assert mock_deliver_alert.call_count == 1
-        assert len(alert2.logs) == 1
-        assert alert2.logs[0].alert_id == 2
-        assert alert2.logs[0].state == "trigger"
+    run_alert_query(database.query(Alert).filter_by(id=2).one().id, database)
+    alert2 = database.query(Alert).filter_by(id=2).one()
+    assert mock_deliver.call_count == 1
+    assert len(alert2.logs) == 1
+    assert alert2.logs[0].alert_id == 2
+    assert alert2.logs[0].state == "trigger"
 
-        run_alert_query(db.session.query(Alert).filter_by(id=3).one().id)
-        alert3 = db.session.query(Alert).filter_by(id=3).one()
-        assert mock_deliver_alert.call_count == 1
-        assert mock_error.call_count == 2
-        assert len(alert3.logs) == 1
-        assert alert3.logs[0].alert_id == 3
-        assert alert3.logs[0].state == "error"
+    run_alert_query(database.query(Alert).filter_by(id=3).one().id, database)
+    alert3 = database.query(Alert).filter_by(id=3).one()
+    assert mock_deliver.call_count == 1
+    assert mock_error.call_count == 2
+    assert len(alert3.logs) == 1
+    assert alert3.logs[0].alert_id == 3
+    assert alert3.logs[0].state == "error"
 
-        run_alert_query(db.session.query(Alert).filter_by(id=4).one().id)
-        assert mock_deliver_alert.call_count == 1
-        assert mock_error.call_count == 3
+    run_alert_query(database.query(Alert).filter_by(id=4).one().id, database)
+    assert mock_deliver.call_count == 1
+    assert mock_error.call_count == 3
 
-        run_alert_query(db.session.query(Alert).filter_by(id=5).one().id)
-        assert mock_deliver_alert.call_count == 1
-        assert mock_error.call_count == 4
+    run_alert_query(database.query(Alert).filter_by(id=5).one().id, database)
+    assert mock_deliver.call_count == 1
+    assert mock_error.call_count == 4
 
 
 @patch("superset.tasks.schedules.deliver_alert")
 @patch("superset.tasks.schedules.run_alert_query")
-def test_schedule_alert_query(mock_run_alert, mock_deliver_alert):
-    with app.app_context():
-        active_alert = db.session.query(Alert).filter_by(id=1).one()
-        inactive_alert = db.session.query(Alert).filter_by(id=3).one()
+def test_schedule_alert_query(mock_run_alert, mock_deliver_alert, setup_database):
+    database = setup_database
+    active_alert = database.query(Alert).filter_by(id=1).one()
+    inactive_alert = database.query(Alert).filter_by(id=3).one()
 
-        # Test that inactive alerts are no processed
-        schedule_alert_query(
-            report_type=ScheduleType.alert, schedule_id=inactive_alert.id
-        )
-        assert mock_run_alert.call_count == 0
-        assert mock_deliver_alert.call_count == 0
+    # Test that inactive alerts are no processed
+    schedule_alert_query(report_type=ScheduleType.alert, schedule_id=inactive_alert.id)
+    assert mock_run_alert.call_count == 0
+    assert mock_deliver_alert.call_count == 0
 
-        # Test that active alerts with no recipients passed in are processed regularly
-        schedule_alert_query(
-            report_type=ScheduleType.alert, schedule_id=active_alert.id
-        )
-        assert mock_run_alert.call_count == 1
-        assert mock_deliver_alert.call_count == 0
+    # Test that active alerts with no recipients passed in are processed regularly
+    schedule_alert_query(report_type=ScheduleType.alert, schedule_id=active_alert.id)
+    assert mock_run_alert.call_count == 1
+    assert mock_deliver_alert.call_count == 0
 
-        # Test that active alerts sent as a test are delivered immediately
-        schedule_alert_query(
-            report_type=ScheduleType.alert,
-            schedule_id=active_alert.id,
-            recipients="testing@email.com",
-            is_test_alert=True,
-        )
-        assert mock_run_alert.call_count == 1
-        assert mock_deliver_alert.call_count == 1
+    # Test that active alerts sent as a test are delivered immediately
+    schedule_alert_query(
+        report_type=ScheduleType.alert,
+        schedule_id=active_alert.id,
+        recipients="testing@email.com",
+        is_test_alert=True,
+    )
+    assert mock_run_alert.call_count == 1
+    assert mock_deliver_alert.call_count == 1