ECSOperator: fix KeyError on missing exitCode (#20264)

* ECSOperator: fix KeyError on missing exitCode

* [tests] add unit test for ECSOperator initialization failure

* check that exit code is not included in the context

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

@@ -470,7 +470,7 @@ def _check_success_task(self) -> None:
                 )
             containers = task['containers']
             for container in containers:
-                if container.get('lastStatus') == 'STOPPED' and container['exitCode'] != 0:
+                if container.get('lastStatus') == 'STOPPED' and container.get('exitCode', 1) != 0:
                     if self.task_log_fetcher:
                         last_logs = "\n".join(
                             self.task_log_fetcher.get_last_log_messages(self.number_logs_exception)

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

@@ -327,6 +327,26 @@ def test_check_success_tasks_raises_logs_disabled(self):
         assert "'exitCode': 1" in str(ctx.value)
         client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
 
+    def test_check_success_tasks_handles_initialization_failure(self):
+        client_mock = mock.Mock()
+        self.ecs.arn = 'arn'
+        self.ecs.client = client_mock
+
+        # exitCode is missing during some container initialization failures
+        client_mock.describe_tasks.return_value = {
+            'tasks': [{'containers': [{'name': 'foo', 'lastStatus': 'STOPPED'}]}]
+        }
+
+        with pytest.raises(Exception) as ctx:
+            self.ecs._check_success_task()
+
+        print(str(ctx.value))
+        assert "This task is not in success state " in str(ctx.value)
+        assert "'name': 'foo'" in str(ctx.value)
+        assert "'lastStatus': 'STOPPED'" in str(ctx.value)
+        assert "exitCode" not in str(ctx.value)
+        client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
+
     def test_check_success_tasks_raises_pending(self):
         client_mock = mock.Mock()
         self.ecs.client = client_mock