always push ECS task ARN to xcom in EcsRunTaskOperator (#33703)

* always push ECS task ARN to xcom in EcsRunTaskOperator

Author: vandonr-amz

airflow/providers/amazon/aws/operators/ecs.py

@@ -543,6 +543,9 @@ def execute(self, context):
             # start the task except if we reattached to an existing one just before.
             self._start_task()
 
+        if self.do_xcom_push:
+            self.xcom_push(context, key="ecs_task_arn", value=self.arn)
+
         if self.deferrable:
             self.defer(
                 trigger=TaskDoneTrigger(

tests/providers/amazon/aws/operators/test_ecs.py

@@ -265,6 +265,7 @@ def test_template_fields_overrides(self):
             ],
         ],
     )
+    @mock.patch.object(EcsRunTaskOperator, "xcom_push")
     @mock.patch.object(EcsRunTaskOperator, "_wait_for_task_ended")
     @mock.patch.object(EcsRunTaskOperator, "_check_success_task")
     @mock.patch.object(EcsBaseOperator, "client")
@@ -273,6 +274,7 @@ def test_execute_without_failures(
         client_mock,
         check_mock,
         wait_mock,
+        xcom_mock,
         launch_type,
         capacity_provider_strategy,
         platform_version,
@@ -626,28 +628,31 @@ def test_reattach_save_task_arn_xcom(
         assert self.ecs.arn == f"arn:aws:ecs:us-east-1:012345678910:task/{TASK_ID}"
         assert "No active previously launched task found to reattach" in caplog.messages
 
+    @mock.patch.object(EcsRunTaskOperator, "xcom_push")
     @mock.patch.object(EcsBaseOperator, "client")
     @mock.patch("airflow.providers.amazon.aws.utils.task_log_fetcher.AwsTaskLogFetcher")
-    def test_execute_xcom_with_log(self, log_fetcher_mock, client_mock):
+    def test_execute_xcom_with_log(self, log_fetcher_mock, client_mock, xcom_mock):
         self.ecs.do_xcom_push = True
         self.ecs.task_log_fetcher = log_fetcher_mock
 
         log_fetcher_mock.get_last_log_message.return_value = "Log output"
 
         assert self.ecs.execute(None) == "Log output"
 
+    @mock.patch.object(EcsRunTaskOperator, "xcom_push")
     @mock.patch.object(EcsBaseOperator, "client")
     @mock.patch("airflow.providers.amazon.aws.utils.task_log_fetcher.AwsTaskLogFetcher")
-    def test_execute_xcom_with_no_log(self, log_fetcher_mock, client_mock):
+    def test_execute_xcom_with_no_log(self, log_fetcher_mock, client_mock, xcom_mock):
         self.ecs.do_xcom_push = True
         self.ecs.task_log_fetcher = log_fetcher_mock
 
         log_fetcher_mock.get_last_log_message.return_value = None
 
         assert self.ecs.execute(None) is None
 
+    @mock.patch.object(EcsRunTaskOperator, "xcom_push")
     @mock.patch.object(EcsBaseOperator, "client")
-    def test_execute_xcom_with_no_log_fetcher(self, client_mock):
+    def test_execute_xcom_with_no_log_fetcher(self, client_mock, xcom_mock):
         self.ecs.do_xcom_push = True
         assert self.ecs.execute(None) is None
 
@@ -657,8 +662,9 @@ def test_execute_xcom_disabled(self, log_fetcher_mock, client_mock):
         self.ecs.do_xcom_push = False
         assert self.ecs.execute(None) is None
 
+    @mock.patch.object(EcsRunTaskOperator, "xcom_push")
     @mock.patch.object(EcsRunTaskOperator, "client")
-    def test_with_defer(self, client_mock):
+    def test_with_defer(self, client_mock, xcom_mock):
         self.ecs.deferrable = True
 
         client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES

tests/system/providers/amazon/aws/example_ecs_fargate.py

@@ -122,8 +122,6 @@ def delete_cluster(cluster_name: str) -> None:
                 "assignPublicIp": "ENABLED",
             },
         },
-        # You must set `reattach=True` in order to get ecs_task_arn if you plan to use a Sensor.
-        reattach=True,
     )
     # [END howto_operator_ecs]