aeroyorch commented on code in PR #65932:
URL: https://github.com/apache/airflow/pull/65932#discussion_r4154942161
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -2832,6 +2833,88 @@ def
test_schedule_tis_up_for_reschedule_does_not_increment_try_number(dag_maker,
assert refreshed_ti.try_number == 3
[email protected]_plugin_manager(plugins=[TestPriorityWeightStrategyPlugin])
+def test_schedule_tis_refreshes_task_instance_only_on_retry(dag_maker,
session):
+ with dag_maker(session=session) as dag:
+ for task_id in ("first_attempt", "retry", "reschedule"):
+ BashOperator(task_id=task_id, bash_command="echo 1",
weight_rule=DecreasingPriorityStrategy())
+
+ dr = dag_maker.create_dagrun(session=session)
+ tis = {ti.task_id: ti for ti in dr.get_task_instances(session=session)}
+ for task_id, state, try_number in (
+ ("first_attempt", None, 0),
+ ("retry", TaskInstanceState.UP_FOR_RETRY, 2),
+ ("reschedule", TaskInstanceState.UP_FOR_RESCHEDULE, 3),
+ ):
+ tis[task_id].refresh_from_task(dag.get_task(task_id))
+ tis[task_id].state = state
+ tis[task_id].try_number = try_number
+ session.commit()
+
+ hook_calls = []
+
+ def route_to_retry_queue(task_instance, dag_run=None):
+ hook_calls.append((task_instance.task_id, task_instance.try_number,
dag_run))
+ task_instance.queue = "retry_queue"
+
+ with _registered_mutation_hook(route_to_retry_queue):
+ assert dr.schedule_tis(tis.values(), session=session) == 3
+ session.commit()
+
+ assert hook_calls == [("retry", 2, dr)]
+ session.expire_all()
+ retry_ti = dr.get_task_instance("retry", session=session)
+ assert retry_ti.state == TaskInstanceState.SCHEDULED
+ assert retry_ti.try_number == 2
+ assert retry_ti.queue == "retry_queue"
+ assert retry_ti.priority_weight == 2
+
+
[email protected]_serialized_dag
+def test_schedule_tis_refreshes_a_retry_that_defers_from_trigger(dag_maker,
session):
+ class TestOperator(BaseOperator):
+ start_trigger_args = StartTriggerArgs(
+ trigger_cls="airflow.triggers.testing.SuccessTrigger",
+ trigger_kwargs=None,
+ next_method="execute_complete",
+ timeout=None,
+ )
+ start_from_trigger = True
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.start_trigger_args.trigger_kwargs = {}
+
+ def execute_complete(self):
+ pass
+
+ with dag_maker(session=session):
+ TestOperator(task_id="task")
+
+ dr = dag_maker.create_dagrun(session=session)
+ ti = dr.get_task_instance("task", session=session)
+ ti.task = dr.dag.get_task("task")
Review Comment:
use `ti = dag_maker.create_ti("task", dag_run=dr)` instead
--
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]
