(airflow) 40/40: [v3-1-test] Fix flaky test_get_last_ti_filters_removed_tasks and Pydantic v2 discriminator error to fix CI (#59273) (#59278)

[ephraimanierobi]

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

ephraimanierobi pushed a commit to branch v3-1-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 6d92f727694711f3bfaf3dd462681d1d41b8b21a
Author: github-actions[bot] 
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Wed Dec 10 16:48:35 2025 +0100

    [v3-1-test] Fix flaky test_get_last_ti_filters_removed_tasks and Pydantic 
v2 discriminator error to fix CI (#59273) (#59278)
    
    * Fix flaky test_get_last_ti_filters_removed_tasks by removing order 
assumptions
    
    * Fix Pydantic v2 callable discriminator error in TIStateUpdate
    (cherry picked from commit 274e096e6df62efcd5d877025d30a23be27c1734)
    
    Co-authored-by: Ankit Chaurasia <8670962+sunank...@users.noreply.github.com>
---
 .../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 f2e7aa36b0d..daa073b5fe5 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
@@ -23,7 +23,6 @@ from typing import Annotated, Any, Literal
 
 from pydantic import (
     AwareDatetime,
-    Discriminator,
     Field,
     JsonValue,
     Tag,
@@ -226,7 +225,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 a3192bc4098..42d6af6bad3 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -2991,11 +2991,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)