Merge pull request #495 from its-mirus-lu/feat/select-all-option

Feature: Implement the ability to select all roles in the role selection screen with "a" or "all"

Author: epierce

gimme_aws_creds/main.py

@@ -420,25 +420,30 @@ def _get_user_int_selections_many(self, min_int, max_int, max_retries=5):
         for _ in range(max_retries):
             selections = set()
             error = False
+            input_values = self.ui.input('Selections (comma separated; specify "A" for all profiles): ').lower().split(',')
 
-            for value in self.ui.input('Selections (comma separated): ').split(','):
-                value = value.strip()
-
-                if not value:
-                    continue
-
-                try:
-                    selection = int(value)
-                except ValueError:
-                    self.ui.warning('Invalid selection {}, must be an integer value.'.format(repr(value)))
-                    error = True
-                    continue
-
-                if min_int <= selection <= max_int:
-                    selections.add(value)
-                else:
-                    self.ui.warning(
-                        'Selection {} out of range <{}, {}>'.format(repr(selection), min_int, max_int))
+            if 'a' in input_values or 'all' in input_values:
+                selections = set(range(min_int, max_int+1))
+                self.ui.message("Generating credentials for profile indices {} to {} ".format(min_int, max_int))
+            else:
+                for value in input_values:
+                    value = value.strip()
+
+                    if not value:
+                        continue
+
+                    try:
+                        selection = int(value)
+                    except ValueError:
+                        self.ui.warning('Invalid selection {}, must be an integer value, or "A" for all profiles'.format(repr(value)))
+                        error = True
+                        continue
+
+                    if min_int <= selection <= max_int:
+                        selections.add(value)
+                    else:
+                        self.ui.warning(
+                            'Selection {} out of range <{}, {}>'.format(repr(selection), min_int, max_int))
 
             if error:
                 continue

tests/test_main.py

@@ -47,12 +47,26 @@ def test_choose_roles_app_2(self, mock):
         self.assertRaises(errors.GimmeAWSCredsExitBase, creds._choose_roles, self.APP_INFO)
         self.assertRaises(errors.GimmeAWSCredsExitBase, creds._choose_app, self.AWS_INFO)
 
-    @patch('builtins.input', return_value='a')
-    def test_choose_roles_app_a(self, mock):
+    @patch('builtins.input', return_value='b')
+    def test_choose_roles_app_b(self, mock):
         creds = GimmeAWSCreds()
         self.assertRaises(errors.GimmeAWSCredsExitBase, creds._choose_roles, self.APP_INFO)
         self.assertRaises(errors.GimmeAWSCredsExitBase, creds._choose_app, self.AWS_INFO)
 
+    @patch('builtins.input', return_value='all')
+    def test_choose_roles_app_select_all(self, mock):
+        creds = GimmeAWSCreds()
+        selections = creds._choose_roles(self.APP_INFO)
+        expected_selections = {self.APP_INFO[0].role, self.APP_INFO[1].role}
+        self.assertEqual(selections, expected_selections)
+        
+    @patch('builtins.input', return_value='a')
+    def test_choose_roles_app_select_all_a(self, mock):
+        creds = GimmeAWSCreds()
+        selections = creds._choose_roles(self.APP_INFO)
+        expected_selections = {self.APP_INFO[0].role, self.APP_INFO[1].role}
+        self.assertEqual(selections, expected_selections)
+        
     def test_get_selected_app_from_config_0(self):
         creds = GimmeAWSCreds()