antonlin1 commented on code in PR #74852:
URL: https://github.com/apache/airflow/pull/74852#discussion_r4443319311
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -13988,6 +13991,74 @@ def _produce_and_register_asset_event(
return apdr
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows")
[email protected]("temporal", [False, True])
[email protected]("event_kind", ["plain", "asset", "partitioned"])
+def test_watcher_event_creates_partitioned_consumer_run(dag_maker, session,
temporal, event_kind):
+ asset = Asset(name="watcher-partitioned-asset")
+ dag_id = "watcher-partitioned-consumer"
+ with dag_maker(
+ dag_id=dag_id,
+ schedule=PartitionedAssetTimetable(
+ assets=asset,
+ default_partition_mapper=StartOfDayMapper() if temporal else
IdentityMapper(),
+ ),
+ session=session,
+ ):
+ EmptyOperator(task_id="consume")
+ ordinary_dag_id = "watcher-ordinary-consumer"
+ with dag_maker(dag_id=ordinary_dag_id, schedule=asset, session=session):
+ EmptyOperator(task_id="consume")
+ session.commit()
+
+ trigger = Trigger(classpath="airflow.triggers.testing.SuccessTrigger",
kwargs={})
+ session.add(trigger)
+ session.flush()
+ asset_model = session.scalar(select(AssetModel).where(AssetModel.name ==
asset.name))
+ asset_model.add_trigger(trigger, "partition-watcher")
+ session.flush()
+ key = "2026-10-16T12:00:00"
+ payload = {"type": "AssetTriggerEvent", "partition_key": key}
+ event: TriggerEvent
+ if event_kind == "plain":
+ event = TriggerEvent(payload)
+ else:
+ event = AssetTriggerEvent(payload, partition_key=key if event_kind ==
"partitioned" else None)
+ adapter = TypeAdapter(ToTriggerSupervisor)
+ changes = messages.TriggerStateChanges(events=[(trigger.id, event, 7)])
+ decoded = adapter.validate_json(adapter.dump_json(changes))
+ Trigger.submit_event(trigger.id, decoded.events[0].event, session=session)
+ session.flush()
+
+ asset_event = session.scalar(select(AssetEvent).where(AssetEvent.asset_id
== asset_model.id))
+ partitioned = event_kind == "partitioned"
+ assert asset_event.partition_key == (key if partitioned else None)
+ assert asset_event.extra == {"from_trigger": True, "payload": payload}
+ ordinary_candidate = session.scalar(
+ select(AssetDagRunQueue).where(AssetDagRunQueue.target_dag_id ==
ordinary_dag_id)
+ )
+ assert (ordinary_candidate is not None) is not partitioned
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ assert runner._create_dagruns_for_partitioned_asset_dags(session=session)
== (
+ {dag_id} if partitioned else set()
+ )
+ runs = session.scalars(select(DagRun).where(DagRun.dag_id == dag_id)).all()
+ assert len(runs) == int(partitioned)
+ if partitioned:
+ [run] = runs
+ assert run.partition_key == ("2026-10-16" if temporal else key)
+ assert run.logical_date is None
+ assert run.state == DagRunState.QUEUED
+ assert [event.id for event in run.consumed_asset_events] ==
[asset_event.id]
+ if temporal:
+ assert run.partition_date == pendulum.datetime(2026, 10, 16,
tz="UTC")
+ else:
+ assert run.partition_date is None
Review Comment:
updated like so. Thank you!
https://github.com/apache/airflow/pull/74852/changes/a724308f53c168170933328b61f75e9fbbf0cdbc#diff-61ce9a0b773f700a900c02608ea61113041507c7a1634eec80e5c416784054b8R14182-R14283
--
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]
