(airflow) branch main updated: Fix flaky test_get_last_ti_filters_removed_tasks and Pydantic v2 discriminator error to fix CI (#59273)

[potiuk]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 274e096e6df Fix flaky test_get_last_ti_filters_removed_tasks and 
Pydantic v2 discriminator error to fix CI (#59273)
274e096e6df is described below

commit 274e096e6df62efcd5d877025d30a23be27c1734
Author: Ankit Chaurasia <8670962+sunank...@users.noreply.github.com>
AuthorDate: Wed Dec 10 21:27:04 2025 +0545

    Fix flaky test_get_last_ti_filters_removed_tasks and Pydantic v2 
discriminator error to fix CI (#59273)
    
    * Fix flaky test_get_last_ti_filters_removed_tasks by removing order 
assumptions
    
    * Fix Pydantic v2 callable discriminator error in TIStateUpdate
---
 .../airflow/api_fastapi/execution_api/datamodels/taskinstance.py | 3 +--
 airflow-core/tests/unit/models/test_dagrun.py                    | 9 ++++++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git 
a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py 
b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
index 56656bc7222..78faf620991 100644
--- 
a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
+++ 
b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
@@ -24,7 +24,6 @@ from typing import Annotated, Any, Literal
 
 from pydantic import (
     AwareDatetime,
-    Discriminator,
     Field,
     JsonValue,
     Tag,
@@ -227,7 +226,7 @@ TIStateUpdate = Annotated[
     | Annotated[TIDeferredStatePayload, Tag("deferred")]
     | Annotated[TIRescheduleStatePayload, Tag("up_for_reschedule")]
     | Annotated[TIRetryStatePayload, Tag("up_for_retry")],
-    Discriminator(ti_state_discriminator),
+    Field(discriminator=ti_state_discriminator),
 ]
 
 
diff --git a/airflow-core/tests/unit/models/test_dagrun.py 
b/airflow-core/tests/unit/models/test_dagrun.py
index ef6b5095fc5..daeff282424 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -2992,11 +2992,14 @@ class TestDagRunGetLastTi:
         dr = dag_maker.create_dagrun()
 
         tis = dr.get_task_instances(session=session)
+        assert len(tis) == 3
+
+        ti_by_id = {ti.task_id: ti for ti in tis}
 
         # Mark some TIs as removed
-        tis[0].state = TaskInstanceState.REMOVED
-        tis[1].state = TaskInstanceState.REMOVED
-        tis[2].state = TaskInstanceState.SUCCESS
+        ti_by_id["task1"].state = TaskInstanceState.REMOVED
+        ti_by_id["task2"].state = TaskInstanceState.REMOVED
+        ti_by_id["task3"].state = TaskInstanceState.SUCCESS
         session.commit()
 
         last_ti = dr.get_last_ti(dag, session=session)