Fix chunk iteration for user migration cmd

keycloak_user_export/management/commands/migrate_users_to_keycloak.py

@@ -8,7 +8,6 @@
 import requests
 
 from keycloak_user_export.models import UserExportToKeycloak
-from open_discussions.utils import queryset_chunks
 
 User = get_user_model()
 
@@ -209,7 +208,6 @@ def _log_api_call(self, response):
         log_out.write(f"Request:\t{response.request.body.decode('utf-8')}")
         log_out.write(f"Response:\t{response.text}")
 
-
     def handle(self, *args, **kwargs):
         self.verbose = kwargs["verbose"]
 
@@ -240,7 +238,9 @@ def handle(self, *args, **kwargs):
 
             # Process batches of the users who must be exported.
             batch_size = kwargs["batch_size"]
-            for batch in queryset_chunks(unsynced_users, chunk_size=batch_size):
+            # the users are updated in such a way that they're excluded from this query the next time
+            # so we just keep grabbing a slice from the beginning of the query until it returns empty
+            while batch := unsynced_users[:batch_size]:
                 payload = {
                     "ifResourceExists": "SKIP",
                     "realm": settings.KEYCLOAK_REALM_NAME,

keycloak_user_export/management/commands/tests/test_migrate_users_to_keycloak.py

@@ -25,8 +25,8 @@ def test_migrate_users_to_keycloak(mocked_responses, django_assert_num_queries):
     """Test that the migrate_users_to_keycloak command functions as expected"""
     num_total_users = 100
     num_already_synced = 20
-    num_to_sync = num_total_users - num_already_synced
     batch_size = 10
+    num_batches = (num_total_users - num_already_synced) / batch_size
     users = UserFactory.create_batch(num_total_users)
     inactive_users = UserFactory.create_batch(10, is_active=False)
     already_synced = users[:num_already_synced]
@@ -73,18 +73,13 @@ def partial_import_callback(request):
         content_type="application/json",
     )
 
-
-
-    with django_assert_num_queries(
-        # 3 per batch - 1 for User, 1 for Profile, 1 for UserExportToKeycloak
-        # plus one for the end of the iteration
-        ceil(num_to_sync / batch_size) * 3 + 1
-    ):
+    # 3 queries per batch plus 1 for the last empty query to terminate the loop
+    with django_assert_num_queries(num_batches * 3 + 1):
         call_command(
             "migrate_users_to_keycloak",
             "--client-id=test-client-id",
             "--client-secret=test-client-secret",
-            f"--batch-size={batch_size}"
+            f"--batch-size={batch_size}",
         )
 
     assert (