1 files changed, 16 insertions, 12 deletions
diff --git a/lib/gitlab/github_import/single_endpoint_notes_importing.rb b/lib/gitlab/github_import/single_endpoint_notes_importing.rb
index 0a3559adde3..aea4059dfbc 100644
--- a/lib/gitlab/github_import/single_endpoint_notes_importing.rb
+++ b/lib/gitlab/github_import/single_endpoint_notes_importing.rb
@@ -63,23 +63,27 @@ module Gitlab
         mark_as_imported(associated)
       end
 
-      def each_associated_page
+      def each_associated_page(&block)
         parent_collection.each_batch(of: BATCH_SIZE, column: :iid) do |batch|
-          batch.each do |parent_record|
-            # The page counter needs to be scoped by parent_record to avoid skipping
-            # pages of notes from already imported parent_record.
-            page_counter = PageCounter.new(project, page_counter_id(parent_record))
-            repo = project.import_source
-            options = collection_options.merge(page: page_counter.current)
+          process_batch(batch, &block)
+        end
+      end
 
-            client.each_page(collection_method, repo, parent_record.iid, options) do |page|
-              next unless page_counter.set(page.number)
+      def process_batch(batch)
+        batch.each do |parent_record|
+          # The page counter needs to be scoped by parent_record to avoid skipping
+          # pages of notes from already imported parent_record.
+          page_counter = PageCounter.new(project, page_counter_id(parent_record))
+          repo = project.import_source
+          options = collection_options.merge(page: page_counter.current)
 
-              yield parent_record, page
-            end
+          client.each_page(collection_method, repo, parent_record.iid, options) do |page|
+            next unless page_counter.set(page.number)
 
-            mark_parent_imported(parent_record)
+            yield parent_record, page
           end
+
+          mark_parent_imported(parent_record)
         end
       end