add UnhealthyNodeReplacement to CLI (#8563)

* add UnhealthyNodeReplacement to CLI

* Undocument set-unhealthy-node-replacement

* Add comment with context and future guidance.

* Add tests for SetUnhealthyNodeReplacement

* PR feedback

awscli/customizations/emr/createcluster.py

@@ -78,6 +78,11 @@ class CreateCluster(Command):
          'help_text': helptext.TERMINATION_PROTECTED},
         {'name': 'no-termination-protected', 'action': 'store_true',
          'group_name': 'termination_protected'},
+        {'name': 'unhealthy-node-replacement', 'action': 'store_true',
+        'group_name': 'unhealthy_node_replacement',
+        'help_text': helptext.UNHEALTHY_NODE_REPLACEMENT},
+        {'name': 'no-unhealthy-node-replacement', 'action': 'store_true',
+        'group_name': 'unhealthy_node_replacement'},
         {'name': 'scale-down-behavior',
          'help_text': helptext.SCALE_DOWN_BEHAVIOR},
         {'name': 'visible-to-all-users', 'action': 'store_true',
@@ -250,6 +255,14 @@ def _run_main_command(self, parsed_args, parsed_globals):
                 '--termination-protected',
                 parsed_args.no_termination_protected,
                 '--no-termination-protected')
+
+        if (parsed_args.unhealthy_node_replacement or parsed_args.no_unhealthy_node_replacement):
+            instances_config['UnhealthyNodeReplacement'] = \
+            emrutils.apply_boolean_options(
+                parsed_args.unhealthy_node_replacement,
+                '--unhealthy-node-replacement',
+                parsed_args.no_unhealthy_node_replacement,
+                '--no-unhealthy-node-replacement')
 
         if (parsed_args.visible_to_all_users is False and
                 parsed_args.no_visible_to_all_users is False):

awscli/customizations/emr/emr.py

@@ -64,3 +64,11 @@ def register_commands(command_table, session, **kwargs):
     command_table['socks'] = ssh.Socks(session)
     command_table['get'] = ssh.Get(session)
     command_table['put'] = ssh.Put(session)
+
+    # Usually, unwanted EMR API commands are removed through awscli/customizations/removals.py at
+    # launch time. However, the SetUnhealthyNodeReplacement API was launched without the removal
+    # in-place. To avoid breaking potential usage, it is undocumented here instead of being removed.
+    # Assuming there is no usage of the 'set-unhealthy-node-replacement' command, it may be moved to
+    # awscli/customizations/removals.py in the future. No future API commands should be undocumented
+    # here; add to awscli/customizations/removals.py at launch time instead.
+    command_table['set-unhealthy-node-replacement']._UNDOCUMENTED = True

awscli/customizations/emr/exceptions.py

@@ -265,7 +265,8 @@ class MissingClusterAttributesError(EmrError):
     fmt = ('aws: error: Must specify one of the following boolean options: '
            '--visible-to-all-users|--no-visible-to-all-users, '
            '--termination-protected|--no-termination-protected, '
-           '--auto-terminate|--no-auto-terminate.')
+           '--auto-terminate|--no-auto-terminate, '
+           '--unhealthy-node-replacement|--no-unhealthy-node-replacement.')
 
 
 class InvalidEmrFsArgumentsError(EmrError):

awscli/customizations/emr/helptext.py

@@ -518,3 +518,7 @@
     'to access the same IAM resources that the step can access. '
     'The execution role can be a cross-account IAM Role.</p> '
 )
+
+UNHEALTHY_NODE_REPLACEMENT = (
+    '<p>Unhealthy node replacement for an Amazon EMR cluster.</p> '
+)

awscli/customizations/emr/modifyclusterattributes.py

@@ -19,8 +19,8 @@
 
 class ModifyClusterAttr(Command):
     NAME = 'modify-cluster-attributes'
-    DESCRIPTION = ("Modifies the cluster attributes 'visible-to-all-users' and"
-                   " 'termination-protected'.")
+    DESCRIPTION = ("Modifies the cluster attributes 'visible-to-all-users', "
+                   " 'termination-protected' and 'unhealthy-node-replacement'.")
     ARG_TABLE = [
         {'name': 'cluster-id', 'required': True,
             'help_text': helptext.CLUSTER_ID},
@@ -42,6 +42,12 @@ class ModifyClusterAttr(Command):
         {'name': 'no-auto-terminate', 'required': False, 'action':
             'store_true', 'group_name': 'auto_terminate',
             'help_text': 'Set cluster auto terminate after completing all the steps on or off'},
+        {'name': 'unhealthy-node-replacement', 'required': False, 'action':
+            'store_true', 'group_name': 'UnhealthyReplacement',
+            'help_text': 'Set Unhealthy Node Replacement on or off'},
+        {'name': 'no-unhealthy-node-replacement', 'required': False, 'action':
+            'store_true', 'group_name': 'UnhealthyReplacement',
+            'help_text': 'Set Unhealthy Node Replacement on or off'},
     ]
 
     def _run_main_command(self, args, parsed_globals):
@@ -58,9 +64,14 @@ def _run_main_command(self, args, parsed_globals):
             raise exceptions.MutualExclusiveOptionError(
                 option1='--auto-terminate',
                 option2='--no-auto-terminate')
+        if (args.unhealthy_node_replacement and args.no_unhealthy_node_replacement):
+            raise exceptions.MutualExclusiveOptionError(
+                option1='--unhealthy-node-replacement',
+                option2='--no-unhealthy-node-replacement')
         if not(args.termination_protected or args.no_termination_protected or
                args.visible_to_all_users or args.no_visible_to_all_users or
-               args.auto_terminate or args.no_auto_terminate):
+               args.auto_terminate or args.no_auto_terminate or
+               args.unhealthy_node_replacement or args.no_unhealthy_node_replacement):
             raise exceptions.MissingClusterAttributesError()
 
         if (args.visible_to_all_users or args.no_visible_to_all_users):
@@ -89,5 +100,14 @@ def _run_main_command(self, args, parsed_globals):
             emrutils.call_and_display_response(self._session,
                                                'SetKeepJobFlowAliveWhenNoSteps',
                                                parameters, parsed_globals)
+
+        if (args.unhealthy_node_replacement or args.no_unhealthy_node_replacement):
+            protected = (args.unhealthy_node_replacement and
+                         not args.no_unhealthy_node_replacement)
+            parameters = {'JobFlowIds': [args.cluster_id],
+                          'UnhealthyNodeReplacement': protected}
+            emrutils.call_and_display_response(self._session,
+                                               'SetUnhealthyNodeReplacement',
+                                               parameters, parsed_globals)
 
         return 0

awscli/examples/emr/create-cluster-synopsis.txt

@@ -12,6 +12,7 @@
      [--additional-info <value>]
      [--ec2-attributes <value>]
      [--termination-protected | --no-termination-protected]
+     [--unhealthy-node-replacement | --no-unhealthy-node-replacement]
      [--scale-down-behavior <value>]
      [--visible-to-all-users | --no-visible-to-all-users]
      [--enable-debugging | --no-enable-debugging]

awscli/examples/emr/describe-cluster.rst

@@ -30,6 +30,7 @@
 		        "ServiceRole": "EMR_DefaultRole", 
 		        "Tags": [], 
 		        "TerminationProtected": true, 
+		        "UnhealthyNodeReplacement": true, 
 		        "ReleaseLabel": "emr-4.0.0", 
 		        "NormalizedInstanceHours": 96, 
 		        "InstanceGroups": [
@@ -126,6 +127,7 @@
                 "ServiceRole": "EMR_DefaultRole",
                 "Tags": [],
                 "TerminationProtected": false,
+                "UnhealthyNodeReplacement": false, 
                 "ReleaseLabel": "emr-5.2.0",
                 "NormalizedInstanceHours": 472,
                 "InstanceCollectionType": "INSTANCE_FLEET",
@@ -200,6 +202,7 @@
 	            "Name": "My Cluster",
 	            "Tags": [],
 	            "TerminationProtected": true,
+	            "UnhealthyNodeReplacement": true, 
 	            "RunningAmiVersion": "2.5.4",
 	            "InstanceGroups": [
 	                {

tests/unit/customizations/emr/test_create_cluster_ami_version.py

@@ -476,6 +476,31 @@ def test_termination_protected_and_no_termination_protected(self):
         result = self.run_cmd(cmd, 255)
         self.assertEqual(expected_error_msg, result[1])
 
+    def test_unhealthy_node_replacement(self):
+        cmd = DEFAULT_CMD + '--unhealthy-node-replacement'
+        result = copy.deepcopy(DEFAULT_RESULT)
+        instances = copy.deepcopy(DEFAULT_INSTANCES)
+        instances['UnhealthyNodeReplacement'] = True
+        result['Instances'] = instances
+        self.assert_params_for_cmd(cmd, result)
+
+    def test_no_unhealthy_node_replacement(self):
+        cmd = DEFAULT_CMD + '--no-unhealthy-node-replacement'
+        result = copy.deepcopy(DEFAULT_RESULT)
+        instances = copy.deepcopy(DEFAULT_INSTANCES)
+        instances['UnhealthyNodeReplacement'] = False
+        result['Instances'] = instances
+        self.assert_params_for_cmd(cmd, result)
+
+    def test_unhealthy_node_replacement_and_no_unhealthy_node_replacement(self):
+        cmd = DEFAULT_CMD + \
+            '--unhealthy-node-replacement --no-unhealthy-node-replacement'
+        expected_error_msg = (
+            '\naws: error: cannot use both --unhealthy-node-replacement'
+            ' and --no-unhealthy-node-replacement options together.\n')
+        result = self.run_cmd(cmd, 255)
+        self.assertEqual(expected_error_msg, result[1])
+
     def test_visible_to_all_users(self):
         cmd = DEFAULT_CMD + '--visible-to-all-users'
         self.assert_params_for_cmd(cmd, DEFAULT_RESULT)

tests/unit/customizations/emr/test_create_cluster_release_label.py

@@ -455,6 +455,31 @@ def test_termination_protected_and_no_termination_protected(self):
         result = self.run_cmd(cmd, 255)
         self.assertEqual(expected_error_msg, result[1])
 
+    def test_unhealthy_node_replacement(self):
+        cmd = DEFAULT_CMD + '--unhealthy-node-replacement'
+        result = copy.deepcopy(DEFAULT_RESULT)
+        instances = copy.deepcopy(DEFAULT_INSTANCES)
+        instances['UnhealthyNodeReplacement'] = True
+        result['Instances'] = instances
+        self.assert_params_for_cmd(cmd, result)
+
+    def test_no_unhealthy_node_replacement(self):
+        cmd = DEFAULT_CMD + '--no-unhealthy-node-replacement'
+        result = copy.deepcopy(DEFAULT_RESULT)
+        instances = copy.deepcopy(DEFAULT_INSTANCES)
+        instances['UnhealthyNodeReplacement'] = False
+        result['Instances'] = instances
+        self.assert_params_for_cmd(cmd, result)
+
+    def test_unhealthy_node_replacement_and_no_unhealthy_node_replacement(self):
+        cmd = DEFAULT_CMD + \
+            '--unhealthy-node-replacement --no-unhealthy-node-replacement'
+        expected_error_msg = (
+            '\naws: error: cannot use both --unhealthy-node-replacement'
+            ' and --no-unhealthy-node-replacement options together.\n')
+        result = self.run_cmd(cmd, 255)
+        self.assertEqual(expected_error_msg, result[1])
+
     def test_visible_to_all_users(self):
         cmd = DEFAULT_CMD + '--visible-to-all-users'
         self.assert_params_for_cmd(cmd, DEFAULT_RESULT)

tests/unit/customizations/emr/test_modify_cluster_attributes.py

@@ -52,6 +52,17 @@ def test_auto_terminate(self):
         args = ' --cluster-id j-ABC123456 --auto-terminate'
         cmdline = self.prefix + args
         result = {'JobFlowIds': ['j-ABC123456'], 'KeepJobFlowAliveWhenNoSteps': False}
+
+    def test_unhealthy_node_replacement(self):
+        args = ' --cluster-id j-ABC123456 --unhealthy-node-replacement'
+        cmdline = self.prefix + args
+        result = {'JobFlowIds': ['j-ABC123456'], 'UnhealthyNodeReplacement': True}
+        self.assert_params_for_cmd(cmdline, result)
+
+    def test_no_unhealthy_node_replacement(self):
+        args = ' --cluster-id j-ABC123456 --no-unhealthy-node-replacement'
+        cmdline = self.prefix + args
+        result = {'JobFlowIds': ['j-ABC123456'], 'UnhealthyNodeReplacement': False}
         self.assert_params_for_cmd(cmdline, result)
 
     def test_visible_to_all_and_no_visible_to_all(self):
@@ -84,33 +95,41 @@ def test_auto_terminate_and_no_auto_terminate(self):
         result = self.run_cmd(cmdline, 255)
         self.assertEqual(expected_error_msg, result[1])
 
-    def test_termination_protected_and_visible_to_all(self):
+    def test_can_set_multiple_attributes(self):
         args = ' --cluster-id j-ABC123456 --termination-protected'\
-               ' --visible-to-all-users'
+               ' --visible-to-all-users --unhealthy-node-replacement'
         cmdline = self.prefix + args
         result_set_termination_protection = {
             'JobFlowIds': ['j-ABC123456'], 'TerminationProtected': True}
         result_set_visible_to_all_users = {
             'JobFlowIds': ['j-ABC123456'], 'VisibleToAllUsers': True}
+        result_set_unhealty_node_replacement = {
+            'JobFlowIds': ['j-ABC123456'], 'UnhealthyNodeReplacement': True}
         self.run_cmd(cmdline)
         self.assertDictEqual(
             self.operations_called[0][1], result_set_visible_to_all_users)
         self.assertDictEqual(
             self.operations_called[1][1], result_set_termination_protection)
+        self.assertDictEqual(
+            self.operations_called[2][1], result_set_unhealty_node_replacement)
 
-    def test_termination_protected_and_no_visible_to_all(self):
+    def test_can_set_multiple_attributes_with_no(self):
         args = ' --cluster-id j-ABC123456 --termination-protected'\
-               ' --no-visible-to-all-users'
+               ' --no-visible-to-all-users --unhealthy-node-replacement'
         cmdline = self.prefix + args
         result_set_termination_protection = {
             'JobFlowIds': ['j-ABC123456'], 'TerminationProtected': True}
         result_set_visible_to_all_users = {
             'JobFlowIds': ['j-ABC123456'], 'VisibleToAllUsers': False}
+        result_set_unhealty_node_replacement = {
+            'JobFlowIds': ['j-ABC123456'], 'UnhealthyNodeReplacement': True}
         self.run_cmd(cmdline)
         self.assertDictEqual(
             self.operations_called[0][1], result_set_visible_to_all_users)
         self.assertDictEqual(
             self.operations_called[1][1], result_set_termination_protection)
+        self.assertDictEqual(
+            self.operations_called[2][1], result_set_unhealty_node_replacement)
 
     def test_at_least_one_option(self):
         args = ' --cluster-id j-ABC123456'
@@ -119,7 +138,8 @@ def test_at_least_one_option(self):
             '\naws: error: Must specify one of the following boolean options: '
             '--visible-to-all-users|--no-visible-to-all-users, '
             '--termination-protected|--no-termination-protected, '
-            '--auto-terminate|--no-auto-terminate.\n')
+            '--auto-terminate|--no-auto-terminate, '
+            '--unhealthy-node-replacement|--no-unhealthy-node-replacement.\n')
         result = self.run_cmd(cmdline, 255)
         self.assertEqual(expected_error_msg, result[1])
 

tests/unit/customizations/emr/test_set_unhealthy_node_replacement.py

@@ -0,0 +1,30 @@
+# Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+from awscli.testutils import BaseAWSHelpOutputTest, BaseAWSCommandParamsTest
+
+
+class TestSetUnhealthyNodeReplacement(BaseAWSCommandParamsTest):
+    prefix = 'emr set-unhealthy-node-replacement'
+
+    def test_unhealthy_node_replacement(self):
+        args = ' --cluster-id j-ABC123456 --unhealthy-node-replacement'
+        cmdline = self.prefix + args
+        result = {'JobFlowIds': ['j-ABC123456'], 'UnhealthyNodeReplacement': True}
+        self.assert_params_for_cmd(cmdline, result)
+
+
+class TestSetUnhealthyNodeReplacementHelp(BaseAWSHelpOutputTest):
+    def test_set_unhealthy_node_replacement_is_undocumented(self):
+        self.driver.main(['emr', 'help'])
+        self.assert_not_contains('set-unhealthy-node-replacement')