haseebmalik18 commented on code in PR #70923:
URL: https://github.com/apache/airflow/pull/70923#discussion_r3696631612


##########
airflow-core/tests/unit/utils/test_db_cleanup.py:
##########
@@ -1343,6 +1360,247 @@ def test_extra_filters_keep_in_flight_rows(self):
             assert seeded[state] not in survivors, f"{state} row should be 
cleaned up"
 
 
+class TestCallbackCleanupConfig:
+    """The callback table is registered under the name it actually has in the 
schema."""
+
+    def test_callback_is_configured_under_its_current_name(self):
+        # The table was renamed from callback_request to callback; a config 
entry naming a
+        # table that does not exist is skipped with a warning, so the rows are 
never purged.
+        assert "callback" in config_dict
+        assert "callback_request" not in config_dict
+
+    def test_removed_tables_are_not_configured(self):
+        assert "sla_miss" not in config_dict
+
+    def test_cleaning_callback_pulls_in_its_dependent_deadline_rows(self):
+        # deadline.callback_id cascades from callback, so deadline has to be 
cleaned
+        # (and archived) first, or those rows would vanish unrecorded.
+        selected, _ = _effective_table_names(table_names=["callback"])
+        assert selected == ["deadline", "callback"]
+
+
[email protected]_test
+class TestCallbackCleanup:
+    """Cleanup must purge finished callbacks without disturbing ones that can 
still run."""
+
+    def setup_method(self):
+        from tests_common.test_utils.db import clear_db_callbacks, 
clear_db_deadline
+
+        clear_db_deadline()
+        clear_db_callbacks()
+        with create_session() as session:
+            for name in _get_archived_table_names(["callback", "deadline"], 
session):
+                session.execute(text(f"DROP TABLE {name}"))
+            session.commit()
+
+    def teardown_method(self):
+        from tests_common.test_utils.db import clear_db_callbacks, 
clear_db_deadline
+
+        clear_db_deadline()
+        clear_db_callbacks()
+        with create_session() as session:
+            for name in _get_archived_table_names(["callback", "deadline"], 
session):
+                session.execute(text(f"DROP TABLE {name}"))
+            session.commit()
+
+    @staticmethod
+    def _add_callback(session, state, created_at):
+        from airflow.executors.workloads.callback import CallbackFetchMethod
+        from airflow.models.callback import Callback
+
+        callback = Callback(priority_weight=1)
+        callback.fetch_method = CallbackFetchMethod.IMPORT_PATH
+        callback.state = state
+        callback.created_at = created_at
+        session.add(callback)
+        session.flush()
+        return callback.id
+
+    @staticmethod
+    def _count_callbacks(session, callback_id):
+        from airflow.models.callback import Callback
+
+        return 
session.scalar(select(func.count()).select_from(Callback).where(Callback.id == 
callback_id))
+
+    @staticmethod
+    def _clean_callbacks(cutoff):
+        with create_session() as session:
+            _cleanup_table(
+                **config_dict["callback"].__dict__,
+                clean_before_timestamp=cutoff,
+                dry_run=False,
+                verbose=False,
+                confirm=False,
+                skip_archive=True,
+                session=session,
+            )
+
+    @pytest.mark.parametrize(
+        ("state", "should_survive"),
+        [
+            ("success", False),
+            ("failed", False),
+            (None, False),
+            ("scheduled", True),
+            ("queued", True),
+            ("running", True),

Review Comment:
   Might be nice to add `pending` here too? It's in `ACTIVE_STATES` alongside 
`queued` and `running` (`models/callback.py:66`), and it's the state a fired 
executor callback sits in while it waits for the executor



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to