From d48e0b3eba935c63d18e580852b5a119459c46ab Mon Sep 17 00:00:00 2001
From: vincbeck <vincbeck@amazon.com>
Date: Tue, 16 Jul 2024 11:58:01 -0400
Subject: [PATCH] Deregister the ECS task definition in the system test
 `example_batch`

---
 tests/system/providers/amazon/aws/example_batch.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/system/providers/amazon/aws/example_batch.py b/tests/system/providers/amazon/aws/example_batch.py
index 77c5d21997364..d47fc683e5e84 100644
--- a/tests/system/providers/amazon/aws/example_batch.py
+++ b/tests/system/providers/amazon/aws/example_batch.py
@@ -25,6 +25,7 @@
 from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.batch import BatchCreateComputeEnvironmentOperator, BatchOperator
+from airflow.providers.amazon.aws.operators.ecs import EcsDeregisterTaskDefinitionOperator
 from airflow.providers.amazon.aws.sensors.batch import (
     BatchComputeEnvironmentSensor,
     BatchJobQueueSensor,
@@ -240,6 +241,12 @@ def delete_job_queue(job_queue_name):
         poke_interval=10,
     )
 
+    deregister_task_definition = EcsDeregisterTaskDefinitionOperator(
+        task_id="deregister_task",
+        task_definition=f"{batch_job_definition_name}:1",
+        trigger_rule=TriggerRule.ALL_DONE,
+    )
+
     log_cleanup = prune_logs(
         [
             # Format: ('log group name', 'log stream prefix')
@@ -271,6 +278,7 @@ def delete_job_queue(job_queue_name):
         wait_for_compute_environment_disabled,
         delete_compute_environment(batch_job_compute_environment_name),
         delete_job_definition(batch_job_definition_name),
+        deregister_task_definition,
         log_cleanup,
     )