make platform version as independent parameter of ECSOperator

apache · Aug 2, 2021 · 4d81f2e · 4d81f2e
1 parent bab7c86
commit 4d81f2e
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 8 deletions.
diff --git a/airflow/providers/amazon/aws/operators/ecs.py b/airflow/providers/amazon/aws/operators/ecs.py
@@ -169,7 +169,7 @@ def __init__(
         group: Optional[str] = None,
         placement_constraints: Optional[list] = None,
         placement_strategy: Optional[list] = None,
-        platform_version: str = 'LATEST',
+        platform_version: Optional[str] = None,
         network_configuration: Optional[dict] = None,
         tags: Optional[dict] = None,
         awslogs_group: Optional[str] = None,
@@ -254,11 +254,10 @@ def _start_task(self, context):
 
         if self.capacity_provider_strategy:
             run_opts['capacityProviderStrategy'] = self.capacity_provider_strategy
-            run_opts['platformVersion'] = self.platform_version
         elif self.launch_type:
             run_opts['launchType'] = self.launch_type
-            if self.launch_type == 'FARGATE':
-                run_opts['platformVersion'] = self.platform_version
+        if self.platform_version is not None:
+            run_opts['platformVersion'] = self.platform_version
         if self.group is not None:
             run_opts['group'] = self.group
         if self.placement_constraints is not None:

diff --git a/tests/providers/amazon/aws/operators/test_ecs.py b/tests/providers/amazon/aws/operators/test_ecs.py
@@ -97,23 +97,38 @@ def test_template_fields_overrides(self):
 
     @parameterized.expand(
         [
-            ['EC2', None, None, {'launchType': 'EC2'}],
-            ['FARGATE', None, None, {'launchType': 'FARGATE', 'platformVersion': 'LATEST'}],
             [
                 'EC2',
                 None,
+                None,
+                None,
+                {'launchType': 'EC2'},
+            ],
+            [
+                'FARGATE',
+                None,
+                'LATEST',
+                None,
+                {'launchType': 'FARGATE', 'platformVersion': 'LATEST'},
+            ],
+            [
+                'EC2',
+                None,
+                None,
                 {'testTagKey': 'testTagValue'},
                 {'launchType': 'EC2', 'tags': [{'key': 'testTagKey', 'value': 'testTagValue'}]},
             ],
             [
                 '',
                 None,
+                None,
                 {'testTagKey': 'testTagValue'},
                 {'tags': [{'key': 'testTagKey', 'value': 'testTagValue'}]},
             ],
             [
                 None,
                 {'capacityProvider': 'FARGATE_SPOT'},
+                'LATEST',
                 None,
                 {
                     'capacityProviderStrategy': {'capacityProvider': 'FARGATE_SPOT'},
@@ -123,6 +138,7 @@ def test_template_fields_overrides(self):
             [
                 'FARGATE',
                 {'capacityProvider': 'FARGATE_SPOT', 'weight': 123, 'base': 123},
+                'LATEST',
                 None,
                 {
                     'capacityProviderStrategy': {
@@ -136,6 +152,7 @@ def test_template_fields_overrides(self):
             [
                 'EC2',
                 {'capacityProvider': 'FARGATE_SPOT'},
+                'LATEST',
                 None,
                 {
                     'capacityProviderStrategy': {'capacityProvider': 'FARGATE_SPOT'},
@@ -147,11 +164,21 @@ def test_template_fields_overrides(self):
     @mock.patch.object(ECSOperator, '_wait_for_task_ended')
     @mock.patch.object(ECSOperator, '_check_success_task')
     def test_execute_without_failures(
-        self, launch_type, capacity_provider_strategy, tags, expected_args, check_mock, wait_mock
+        self,
+        launch_type,
+        capacity_provider_strategy,
+        platform_version,
+        tags,
+        expected_args,
+        check_mock,
+        wait_mock,
     ):
 
         self.set_up_operator(
-            launch_type=launch_type, capacity_provider_strategy=capacity_provider_strategy, tags=tags
+            launch_type=launch_type,
+            capacity_provider_strategy=capacity_provider_strategy,
+            platform_version=platform_version,
+            tags=tags,
         )
         client_mock = self.aws_hook_mock.return_value.get_conn.return_value
         client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES