[GitHub] [airflow] feluelle commented on a change in pull request #12952: Add support for capacityProviderStrategy option in ECSOperator.

[GitBox <gi...@apache.org>]

feluelle commented on a change in pull request #12952:
URL: https://github.com/apache/airflow/pull/12952#discussion_r559350937



##########
File path: airflow/providers/amazon/aws/operators/ecs.py
##########
@@ -207,7 +213,9 @@ def _start_task(self):
             'startedBy': self.owner,
         }
 
-        if self.launch_type:
+        if self.capacity_provider_strategy is not None:

Review comment:
       ```suggestion
           if self.capacity_provider_strategy:
   ```

##########
File path: tests/providers/amazon/aws/operators/test_ecs.py
##########
@@ -94,28 +94,50 @@ def test_template_fields_overrides(self):
 
     @parameterized.expand(
         [
-            ['EC2', None],
-            ['FARGATE', None],
-            ['EC2', {'testTagKey': 'testTagValue'}],
-            ['', {'testTagKey': 'testTagValue'}],
+            ['EC2', None, None],
+            ['FARGATE', None, None],
+            ['EC2', {'testTagKey': 'testTagValue'}, None],
+            ['', {'testTagKey': 'testTagValue'}, None],
+            [None, None, [{'capacityProvider': 'FARGATE_SPOT', 'weight': 123, 'base': 123}]],
+            [
+                None,
+                {'testTagKey': 'testTagValue'},
+                [{'capacityProvider': 'FARGATE_SPOT', 'weight': 123, 'base': 123}],
+            ],
+            [
+                'FARGATE',
+                {'testTagKey': 'testTagValue'},
+                [{'capacityProvider': 'FARGATE_SPOT', 'weight': 123, 'base': 123}],
+            ],
+            [
+                'EC2',
+                {'testTagKey': 'testTagValue'},
+                [{'capacityProvider': 'FARGATE_SPOT', 'weight': 123, 'base': 123}],
+            ],
         ]
     )
     @mock.patch.object(ECSOperator, '_wait_for_task_ended')
     @mock.patch.object(ECSOperator, '_check_success_task')
-    def test_execute_without_failures(self, launch_type, tags, check_mock, wait_mock):
+    def test_execute_without_failures(
+        self, launch_type, tags, capacity_provider_strategy, check_mock, wait_mock
+    ):
 
-        self.set_up_operator(launch_type=launch_type, tags=tags)  # pylint: disable=no-value-for-parameter
+        self.set_up_operator(  # pylint: disable=no-value-for-parameter
+            launch_type=launch_type, tags=tags, capacity_provider_strategy=capacity_provider_strategy
+        )
         client_mock = self.aws_hook_mock.return_value.get_conn.return_value
         client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
 
         self.ecs.execute(None)
 
         self.aws_hook_mock.return_value.get_conn.assert_called_once()
         extend_args = {}
-        if launch_type:
+        if capacity_provider_strategy is not None:
+            extend_args['capacityProviderStrategy'] = capacity_provider_strategy
+        elif launch_type:
             extend_args['launchType'] = launch_type
-        if launch_type == 'FARGATE':
-            extend_args['platformVersion'] = 'LATEST'
+            if launch_type == 'FARGATE':
+                extend_args['platformVersion'] = 'LATEST'
         if tags:
             extend_args['tags'] = [{'key': k, 'value': v} for (k, v) in tags.items()]

Review comment:
       Can we pass extended args as a parameter to `@parameterized.expand`? i.e. so that we define what extended args we expect - instead of defining the **same** `if, elif, else` _business_ logic in the tests.
   
   A test generally has input and output. It should not reproduce the exact behaviour of the code you want to test.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org