amoghrajesh commented on code in PR #73030:
URL: https://github.com/apache/airflow/pull/73030#discussion_r4070464984


##########
task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py:
##########
@@ -470,3 +470,37 @@ def test_head_version_keeps_arg_bindings(self, 
real_migrator, startup_details):
         assert isinstance(defaulted, LiteralArgBinding)
         assert defaulted.from_default is True
         assert defaulted.value_schema.root == {"type": "integer", "format": 
"int64"}
+
+
+class TestRealBundleRetryReasonUpgrade:
+    """
+    Drive the *real* supervisor bundle through the ``retry_reason`` migration.
+
+    ``TaskState`` flows foreign-runtime -> supervisor, the opposite direction 
from
+    ``arg_bindings`` above, so a runtime pinned to an older schema is exercised
+    through ``upgrade`` rather than ``downgrade``.
+    """
+
+    @pytest.fixture
+    def real_migrator(self) -> SchemaVersionMigrator:
+        return get_schema_version_migrator()
+
+    def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator):
+        from airflow.sdk.execution_time.comms import TaskState
+
+        body = {"type": "TaskState", "state": "failed", "end_date": None, 
"rendered_map_index": None}
+        out = real_migrator.upgrade(body, TaskState, "2026-06-16")
+        assert out["retry_reason"] is None
+
+    def test_upgrade_keeps_retry_reason_at_head(self, real_migrator):
+        from airflow.sdk.execution_time.comms import TaskState

Review Comment:
   You're right, neither test touched the direction the gate governs. Added a 
downgrade test: retry_reason is gone at `2026-06-16` and kept at `2026-10-30`.
   
   Checked it does what the old ones didn't. Deleting 
`AddRetryReasonToTaskState` from the bundle now fails that test, while the 
others stay green exactly as you measured.
   
   Dropped `test_upgrade_keeps_retry_reason_at_head` as you pointed out it 
passes the supervisor's own version, so it short circuits and only proves 
pydantic round-trips. Kept the upgrade test that backfills None, since that's 
the real foreign-runtime path.



-- 
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]

Reply via email to