agent/authentic2: batch the provisionning of agents (#52620)

hobo/agent/authentic2/provisionning.py

@@ -1,3 +1,4 @@
+from itertools import chain, islice
 import json
 from django.utils.six.moves.urllib.parse import urljoin
 import threading
@@ -26,6 +27,15 @@ RoleParenting = get_role_parenting_model()
 logger = logging.getLogger(__name__)
 
 
+def batch(iterable, size):
+    """Batch an iterable as an iterable of iterables of at most size element
+    long.
+    """
+    sourceiter = iter(iterable)
+    for first in sourceiter:
+        yield chain([first], islice(sourceiter, size - 1))
+
+
 class Provisionning(threading.local):
     __slots__ = ['threads']
 
@@ -166,8 +176,10 @@ class Provisionning(threading.local):
             if roles_with_attributes:
                 for ou, users in ous.items():
                     for service, audience in self.get_audience(ou):
-                        for user in users:
-                            logger.info(u'provisionning user %s to %s', user, audience)
+                        for batched_users in batch(users, 500):
+                            batched_users = list(batched_users)
+                            for user in batched_users:
+                                logger.info('provisionning user %s to %s', user, audience)
                             self.notify_agents({
                                 '@type': 'provision',
                                 'issuer': issuer,
@@ -175,7 +187,7 @@ class Provisionning(threading.local):
                                 'full': False,
                                 'objects': {
                                     '@type': 'user',
-                                    'data': [user_to_json(ou, service, user, user_roles)],
+                                    'data': [user_to_json(ou, service, user, user_roles) for user in batched_users],
                                 }
                             })
             else:

tests_authentic/test_provisionning.py

@@ -278,41 +278,40 @@ def test_provision_user(transactional_db, tenant, caplog):
                 role.members.add(user1)
                 user2.save()
 
-            assert notify_agents.call_count == 3
+            assert notify_agents.call_count == 2
             assert notify_agents.call_args_list[0][0][0]['objects']['@type'] == 'role'
-            for arg in notify_agents.call_args_list[1:]:
-                assert arg == call(ANY)
-                arg = arg[0][0]
-                assert isinstance(arg, dict)
-                assert set(arg.keys()) == set([
-                    'issuer', 'audience', '@type', 'objects', 'full'])
-                assert arg['issuer'] == \
-                    'http://%s/idp/saml2/metadata' % tenant.domain_url
-                assert arg['audience'] == ['http://provider.com']
-                assert arg['@type'] == 'provision'
-                assert arg['full'] is False
-                objects = arg['objects']
-                assert isinstance(objects, dict)
-                assert set(objects.keys()) == set(['data', '@type'])
-                assert objects['@type'] == 'user'
-                data = objects['data']
-                assert isinstance(data, list)
-                assert len(data) == 1
-                for o in data:
-                    assert set(o.keys()) >= set(['uuid', 'username', 'first_name',
-                                                 'is_superuser', 'last_name', 'email', 'roles'])
-                    assert o['uuid'] in users
-                    user = users[o['uuid']]
-                    assert o['uuid'] == user.uuid
-                    assert o['username'] == user.username
-                    assert o['first_name'] == user.first_name
-                    assert o['last_name'] == user.last_name
-                    assert o['email'] == user.email
-                    assert o['roles'] == [{'name': r.name, 'slug': r.slug, 'uuid': r.uuid} for r in
-                                          user.roles.all()]
-                    assert o['is_superuser'] is (user == user1)
-            assert len(set(arg[0][0]['objects']['data'][0]['uuid'] for arg in
-                           notify_agents.call_args_list[1:])) == 2
+            arg = notify_agents.call_args_list[1]
+            assert arg == call(ANY)
+            arg = arg[0][0]
+            assert isinstance(arg, dict)
+            assert set(arg.keys()) == set([
+                'issuer', 'audience', '@type', 'objects', 'full'])
+            assert arg['issuer'] == \
+                'http://%s/idp/saml2/metadata' % tenant.domain_url
+            assert arg['audience'] == ['http://provider.com']
+            assert arg['@type'] == 'provision'
+            assert arg['full'] is False
+            objects = arg['objects']
+            assert isinstance(objects, dict)
+            assert set(objects.keys()) == set(['data', '@type'])
+            assert objects['@type'] == 'user'
+            data = objects['data']
+            assert isinstance(data, list)
+            assert len(data) == 2
+            for o in data:
+                assert set(o.keys()) >= set(['uuid', 'username', 'first_name',
+                                             'is_superuser', 'last_name', 'email', 'roles'])
+                assert o['uuid'] in users
+                user = users[o['uuid']]
+                assert o['uuid'] == user.uuid
+                assert o['username'] == user.username
+                assert o['first_name'] == user.first_name
+                assert o['last_name'] == user.last_name
+                assert o['email'] == user.email
+                assert o['roles'] == [{'name': r.name, 'slug': r.slug, 'uuid': r.uuid} for r in
+                                      user.roles.all()]
+                assert o['is_superuser'] is (user == user1)
+            assert len(set(x['uuid'] for x in notify_agents.call_args_list[1][0][0]['objects']['data'])) == 2
 
             notify_agents.reset_mock()
             with provisionning: