(bugfix): EcsRunTaskOperator decouple volume_configurations from …

…`capacity_provider_strategy` (apache#43047) * Decouple volume_configurations from capacity_provider_strategy * Add tests for volume configurations * Add EOF
astronomer · Nov 6, 2024 · 995cd8f · 995cd8f
1 parent 43f3c31
commit 995cd8f
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 3 deletions.
diff --git a/providers/src/airflow/providers/amazon/aws/operators/ecs.py b/providers/src/airflow/providers/amazon/aws/operators/ecs.py
@@ -598,10 +598,10 @@ def _start_task(self):
 
         if self.capacity_provider_strategy:
             run_opts["capacityProviderStrategy"] = self.capacity_provider_strategy
-            if self.volume_configurations is not None:
-                run_opts["volumeConfigurations"] = self.volume_configurations
         elif self.launch_type:
             run_opts["launchType"] = self.launch_type
+        if self.volume_configurations is not None:
+            run_opts["volumeConfigurations"] = self.volume_configurations
         if self.platform_version is not None:
             run_opts["platformVersion"] = self.platform_version
         if self.group is not None:

diff --git a/providers/tests/amazon/aws/operators/test_ecs.py b/providers/tests/amazon/aws/operators/test_ecs.py
@@ -192,48 +192,54 @@ def test_template_fields_overrides(self):
         )
 
     @pytest.mark.parametrize(
-        "launch_type, capacity_provider_strategy,platform_version,tags,expected_args",
+        "launch_type, capacity_provider_strategy,platform_version,tags,volume_configurations,expected_args",
         [
             [
                 "EC2",
                 None,
                 None,
                 None,
+                None,
                 {"launchType": "EC2"},
             ],
             [
                 "EXTERNAL",
                 None,
                 None,
                 None,
+                None,
                 {"launchType": "EXTERNAL"},
             ],
             [
                 "FARGATE",
                 None,
                 "LATEST",
                 None,
+                None,
                 {"launchType": "FARGATE", "platformVersion": "LATEST"},
             ],
             [
                 "EC2",
                 None,
                 None,
                 {"testTagKey": "testTagValue"},
+                None,
                 {"launchType": "EC2", "tags": [{"key": "testTagKey", "value": "testTagValue"}]},
             ],
             [
                 "",
                 None,
                 None,
                 {"testTagKey": "testTagValue"},
+                None,
                 {"tags": [{"key": "testTagKey", "value": "testTagValue"}]},
             ],
             [
                 None,
                 {"capacityProvider": "FARGATE_SPOT"},
                 "LATEST",
                 None,
+                None,
                 {
                     "capacityProviderStrategy": {"capacityProvider": "FARGATE_SPOT"},
                     "platformVersion": "LATEST",
@@ -244,6 +250,7 @@ def test_template_fields_overrides(self):
                 {"capacityProvider": "FARGATE_SPOT", "weight": 123, "base": 123},
                 "LATEST",
                 None,
+                None,
                 {
                     "capacityProviderStrategy": {
                         "capacityProvider": "FARGATE_SPOT",
@@ -258,11 +265,70 @@ def test_template_fields_overrides(self):
                 {"capacityProvider": "FARGATE_SPOT"},
                 "LATEST",
                 None,
+                None,
                 {
                     "capacityProviderStrategy": {"capacityProvider": "FARGATE_SPOT"},
                     "platformVersion": "LATEST",
                 },
             ],
+            [
+                "FARGATE",
+                None,
+                None,
+                None,
+                [
+                    {
+                        "name": "ebs-volume",
+                        "managedEBSVolume": {
+                            "volumeType": "gp3",
+                            "sizeInGiB": 10,
+                        },
+                        "roleArn": "arn:aws:iam:1111222333:role/ecsInfrastructureRole",
+                    }
+                ],
+                {
+                    "launchType": "FARGATE",
+                    "volumeConfigurations": [
+                        {
+                            "name": "ebs-volume",
+                            "managedEBSVolume": {
+                                "volumeType": "gp3",
+                                "sizeInGiB": 10,
+                            },
+                            "roleArn": "arn:aws:iam:1111222333:role/ecsInfrastructureRole",
+                        }
+                    ],
+                },
+            ],
+            [
+                None,
+                {"capacityProvider": "FARGATE_SPOT"},
+                None,
+                None,
+                [
+                    {
+                        "name": "ebs-volume",
+                        "managedEBSVolume": {
+                            "volumeType": "gp3",
+                            "sizeInGiB": 10,
+                        },
+                        "roleArn": "arn:aws:iam:1111222333:role/ecsInfrastructureRole",
+                    }
+                ],
+                {
+                    "capacityProviderStrategy": {"capacityProvider": "FARGATE_SPOT"},
+                    "volumeConfigurations": [
+                        {
+                            "name": "ebs-volume",
+                            "managedEBSVolume": {
+                                "volumeType": "gp3",
+                                "sizeInGiB": 10,
+                            },
+                            "roleArn": "arn:aws:iam:1111222333:role/ecsInfrastructureRole",
+                        }
+                    ],
+                },
+            ],
         ],
     )
     @mock.patch.object(EcsRunTaskOperator, "xcom_push")
@@ -279,13 +345,15 @@ def test_execute_without_failures(
         capacity_provider_strategy,
         platform_version,
         tags,
+        volume_configurations,
         expected_args,
     ):
         self.set_up_operator(
             launch_type=launch_type,
             capacity_provider_strategy=capacity_provider_strategy,
             platform_version=platform_version,
             tags=tags,
+            volume_configurations=volume_configurations,
         )
         client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES