Refactor group list split logic, add more error and type checking

Author: adorton-adobe

tests/test_connections.py

@@ -369,3 +369,53 @@ def test_no_group_split():
     user.add_to_groups(groups=add_groups, group_type=GroupTypes.usergroup)
     assert user.maybe_split_groups(10) is False
     assert len(user.commands) == 1
+
+
+def test_complex_group_split():
+    """
+    Test a complex command with add and remove directive, with multiple group types
+    UserAction's interface doesn't support this, so we build our own command array
+    :return:
+    """
+    group_prefix = "G"
+    add_groups = [group_prefix+six.text_type(n+1) for n in range(0, 150)]
+    add_products = [group_prefix+six.text_type(n+1) for n in range(0, 26)]
+    user = UserAction(id_type=IdentityTypes.enterpriseID, email="[email protected]")
+    user.commands = [{
+        "add": {
+            GroupTypes.usergroup.name: add_groups,
+            GroupTypes.product.name: add_products,
+        },
+    }]
+    assert user.maybe_split_groups(10) is True
+    assert len(user.commands) == 15
+    assert len([c for c in user.commands if 'product' in c['add']]) == 3
+    assert GroupTypes.product.name not in user.commands[3]['add']
+
+
+def test_split_remove_all():
+    """
+    Don't split groups if "remove" is "all" instead of list
+    :return:
+    """
+    group_prefix = "G"
+    add_groups = [group_prefix+six.text_type(n+1) for n in range(0, 11)]
+    user = UserAction(id_type=IdentityTypes.enterpriseID, email="[email protected]")
+    user.remove_from_groups(all_groups=True)
+    assert user.maybe_split_groups(1) is False
+    assert user.wire_dict() == {'do': [{'remove': 'all'}], 'user': '[email protected]'}
+    user.add_to_groups(groups=add_groups)
+    assert user.maybe_split_groups(10) is True
+    assert user.wire_dict() == {'do': [{'remove': 'all'},
+                                       {'add': {'product': ['G1',
+                                                            'G2',
+                                                            'G3',
+                                                            'G4',
+                                                            'G5',
+                                                            'G6',
+                                                            'G7',
+                                                            'G8',
+                                                            'G9',
+                                                            'G10']}},
+                                       {'add': {'product': ['G11']}}],
+                                'user': '[email protected]'}

umapi_client/functional.py

@@ -43,17 +43,6 @@ class RoleTypes(Enum):
     productAdmin = 2
 
 
-class StepKeys(Enum):
-    add = 1
-    addAdobeID = 2
-    addRoles = 3
-    createEnterpriseID = 4
-    remove = 5
-    removeFromOrg = 6
-    removeRoles = 7
-    update = 8
-
-
 class IfAlreadyExistsOptions(Enum):
     ignoreIfAlreadyExists = 1
     updateIfAlreadyExists = 2
@@ -328,40 +317,32 @@ def maybe_split_groups(self, max_groups):
         new_commands = []
         # return True if we split at least once
         maybe_split = False
+        valid_step_keys = ['add', 'addRoles', 'remove']
         for command in self.commands:
+            # commands are assumed to contain a single key
             step_key, step_args = next(six.iteritems(command))
-            if step_key not in [StepKeys.add.name, StepKeys.remove.name, StepKeys.addRoles.name]:
+            if step_key not in valid_step_keys:
                 new_commands.append(command)
                 continue
-            split_commands = self._split_groups(step_key, command, max_groups)
-            if len(split_commands) > 1:
-                maybe_split = True
+            split_commands = [command]
+            while True:
+                new_command = {step_key: {}}
+                if not isinstance(command[step_key], dict):
+                    break
+                for group_type, groups in six.iteritems(command[step_key]):
+                    if len(groups) > max_groups:
+                        maybe_split = True
+                        command[step_key][group_type], new_command[step_key][group_type] = \
+                            groups[0:max_groups], groups[max_groups:]
+                if new_command[step_key]:
+                    split_commands.append(new_command)
+                    command = new_command
+                else:
+                    break
             new_commands += split_commands
         self.commands = new_commands
         return maybe_split
 
-    def _split_groups(self, step_key, command, max_groups):
-        """
-        Split long group lists of individual command
-        Creates additional commands for each split and recursively splits new groups until no commands exceed group
-        limit.  Assumes that at least one split should occur.
-        :param step_key: Step key (add, remove, etc)
-        :param command: Single command containing one or more group lists to add, remove, etc
-        :param max_groups: Max group list size
-        :return: List of at least two commands that resulted from split
-        """
-        new_commands = [command]
-        while True:
-            new_command = {step_key: {}}
-            for group_type, groups in six.iteritems(command[step_key]):
-                if len(groups) > max_groups:
-                    command[step_key][group_type], new_command[step_key][group_type] = \
-                        groups[0:max_groups], groups[max_groups:]
-            if new_command[step_key]:
-                new_commands += self._split_groups(step_key, new_command, max_groups)
-            else:
-                return new_commands
-
 
 class UsersQuery(QueryMultiple):
     """