hkc-8010 commented on code in PR #66854:
URL: https://github.com/apache/airflow/pull/66854#discussion_r3639433167
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -5712,6 +5713,150 @@ def dict_from_obj(obj):
assert created_run.creating_job_id == scheduler_job.id
+ def _seed_asset_producer_and_queue_row(self, session, dag_maker, *,
outlet_events=None):
+ """Create a producer DAG whose task has an asset outlet, a downstream
DAG scheduled on that
+ asset, a successful producer TI, and a matching ``AssetEventQueue``
row. Returns the TI."""
+ asset = Asset(uri="test://drain-asset", name="drain_asset",
group="test_group")
+
+ with dag_maker(dag_id="drain-producer", start_date=timezone.utcnow(),
session=session) as dag:
+ BashOperator(task_id="task", bash_command="echo 1",
outlets=[asset])
+ dr = dag_maker.create_dagrun(run_id="producer-run")
+ [ti] = dr.get_task_instances(session=session)
+ ti.state = State.SUCCESS
+ # Capture the outlet profiles before creating other DAGs staled the
serialized reference.
+ task_outlets = [o.asprofile().model_dump(mode="json") for o in
dag.get_task("task").outlets]
+ session.flush()
+
+ # Downstream DAG scheduled on the asset -> registration should enqueue
it (ADRQ row).
+ with dag_maker(dag_id="drain-consumer", schedule=[asset],
session=session):
+ EmptyOperator(task_id="downstream")
+ session.add(
+ AssetEventQueue(
+ ti_id=ti.id,
+ payload={"task_outlets": task_outlets, "outlet_events":
outlet_events or []},
+ attempts=0,
+ )
+ )
+ session.commit()
+ return ti
+
+ @pytest.mark.need_serialized_dag
+ def test_drain_asset_event_queue_registers_and_deletes(self, session,
dag_maker):
+ """A pending queue row for a successful producer TI is registered
(AssetEvent + downstream
+ AssetDagRunQueue), the queue row is deleted, and the processed metric
is emitted."""
+ ti = self._seed_asset_producer_and_queue_row(session, dag_maker)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[self.null_exec])
+
+ with mock.patch("airflow.jobs.scheduler_job_runner.stats") as
mock_stats:
+ self.job_runner._drain_asset_event_queue(session=session)
+
+ # The queue row was consumed.
+ assert session.scalars(select(AssetEventQueue)).all() == []
+ # An AssetEvent was produced by the (now drained) TI.
+ event =
session.scalar(select(AssetEvent).where(AssetEvent.source_run_id == ti.run_id))
+ assert event is not None
+ # The downstream asset-scheduled DAG was queued.
+ adrq = session.scalars(
+ select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id ==
"drain-consumer")
+ ).all()
+ assert len(adrq) == 1
+ mock_stats.incr.assert_any_call("asset.event_queue.processed", 1)
+
+ @pytest.mark.need_serialized_dag
+ def test_drain_asset_event_queue_retries_then_parks(self, session,
dag_maker):
+ """A row whose registration keeps failing is retried across drain
passes, one attempt per
+ pass, and once it hits max_attempts it is parked and the failures
metric is emitted."""
+ ti = self._seed_asset_producer_and_queue_row(session, dag_maker)
+
+ scheduler_job = Job()
+ self.job_runner = SchedulerJobRunner(job=scheduler_job,
executors=[self.null_exec])
+
+ max_attempts = self.job_runner._asset_event_queue_max_attempts
+
+ with mock.patch(
+
"airflow.models.taskinstance.TaskInstance.register_asset_changes_in_db",
+ side_effect=RuntimeError("boom"),
+ ):
+ # First pass: the single pending row is claimed once, its
registration fails, so attempts
+ # becomes 1 and it is not yet parked.
+ with mock.patch("airflow.jobs.scheduler_job_runner.stats") as
mock_stats:
+ self.job_runner._drain_asset_event_queue(session=session)
+ row = session.scalars(select(AssetEventQueue)).one()
+ assert row.ti_id == ti.id
+ assert row.attempts == 1
+ assert mock.call("asset.event_queue.failures") not in
mock_stats.incr.mock_calls
+
+ # Run further passes; on the pass that pushes attempts to
max_attempts the row is parked
+ # and the failures metric fires.
+ for _ in range(max_attempts - 1):
Review Comment:
Done. `test_drain_asset_event_queue_retries_then_parks` now loops
`range(max_attempts + 1)`, asserts `attempts == min(cur_attempt + 1,
max_attempts)` on each pass, checks the `failures` metric fires on the pass
that reaches the cap, and the extra pass proves the counter caps (the drain
stops claiming the row once `attempts == max_attempts`).
--
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]
