refactored the specs to the background migration

lib/gitlab/background_migration/migrate_uploads_to_object_storage.rb

+# frozen_string_literal: true
+# rubocop:disable Metrics/LineLength
+# rubocop:disable Style/Documentation
+
 module Gitlab
   module BackgroundMigration
     class MigrateUploadsToObjectStorage
@@ -103,8 +107,8 @@ module Gitlab
         def report(results)
           success, failures = Gitlab::Utils::BisectEnumerable.bisect(results, &:success?)
 
-          puts header(success, failures)
-          puts failures(failures)
+          Rails.logger.info header(success, failures)
+          Rails.logger.warn failures(failures)
         end
 
         def header(success, failures)
@@ -124,7 +128,6 @@ module Gitlab
         sanity_check!(uploads, mounted_as)
 
         BackgroundMigrationWorker.perform_async('MigrateUploadsToObjectStorage', [uploads.ids, mounted_as, to_store])
-        #BackgroundMigration.perform('MigrateUploadsToObjectStorage', [uploads.map(&:id), mounted_as, to_store])
       end
 
       # We need to be sure all the uploads are for the same uploader and model type
@@ -146,7 +149,7 @@ module Gitlab
       end
 
       def perform(ids, mounted_as, to_store)
-        @mounted_as = mounted_as
+        @mounted_as = mounted_as.to_sym
         @to_store = to_store
 
         uploads = Upload.preload(:model).where(id: ids)
@@ -155,9 +158,9 @@ module Gitlab
         results = migrate(uploads)
 
         report(results)
-      rescue SanityCheckError
+      rescue SanityCheckError => e
         # do not retry if the job is insane
-        logger.warn "UploadsToObjectStorage sanity check error: #{e.message}"
+        Rails.logger.warn "UploadsToObjectStorage sanity check error: #{e.message}"
       end
 
       def sanity_check!(uploads)
@@ -168,16 +171,6 @@ module Gitlab
         uploads.map { |upload| upload.build_uploader(@mounted_as) }
       end
 
-      def migrate(batch_size, &block)
-        each_upload_batch(batch_size) do |batch|
-          results = build_uploaders(batch).map(&method(:process_uploader))
-
-          yield results # yield processed batch as [MigrationResult]
-
-          @results.concat(results)
-        end
-      end
-
       def migrate(uploads)
         build_uploaders(uploads).map(&method(:process_uploader))
       end

lib/tasks/gitlab/uploads/migrate.rake

 require_relative 'helpers'
 
-include UploadTaskHelpers
-
 namespace :gitlab do
   namespace :uploads do
     desc 'GitLab | Uploads | Migrate the uploaded files to object storage'
     task :migrate, [:uploader_class, :model_class, :mounted_as] => :environment do |task, args|
-      to_store       = ObjectStorage::Store::REMOTE
+      include UploadTaskHelpers
+
+      @to_store      = ObjectStorage::Store::REMOTE
+      @mounted_as    = args.mounted_as&.gsub(':', '')&.to_sym
       uploader_class = args.uploader_class.constantize
       model_class    = args.model_class.constantize
-      mounted_as     = args.mounted_as&.gsub(':', '')&.to_sym
 
       Upload
-        .where.not(store: to_store)
+        .where.not(store: @to_store)
         .where(uploader: uploader_class.to_s,
                model_type: model_class.to_s)
-        .in_batches(of: batch_size) do |batch| # rubocop: disable Cop/InBatches
-        job = Gitlab::BackgroundMigration::MigrateUploadsToObjectStorage.enqueue!(batch, mounted_as, to_store)
-        puts "Enqueued job: #{job}"
-      end
+        .in_batches(of: batch_size, &method(:process)) # rubocop: disable Cop/InBatches
+    end
+
+    def process(batch)
+      job = Gitlab::BackgroundMigration::MigrateUploadsToObjectStorage.enqueue!(batch,
+                                                                                @mounted_as,
+                                                                                @to_store)
+      puts "Enqueued job: #{job}"
+    rescue Gitlab::BackgroundMigration::MigrateUploadsToObjectStorage::SanityCheckError => e
+      # continue for the next batch
+      puts "Could not enqueue batch (#{batch.ids}) #{e.message}".color(:red)
     end
   end
 end

spec/lib/gitlab/background_migration/migrate_uploads_to_object_storage_spec.rb

@@ -7,15 +7,13 @@ describe Gitlab::BackgroundMigration::MigrateUploadsToObjectStorage, :sidekiq do
     end
   end
 
+  let!(:projects) { create_list(:project, 10, :with_avatar) }
   let(:uploads) { Upload.all }
   let(:mounted_as) { :avatar }
   let(:to_store) { ObjectStorage::Store::REMOTE }
 
   before do
-    stub_env('BATCH', 1)
-    stub_licensed_features(object_storage: true)
-
-    create_list(:upload, 5)
+    stub_uploads_object_storage(AvatarUploader)
   end
 
   describe '.enqueue!' do
@@ -24,6 +22,7 @@ describe Gitlab::BackgroundMigration::MigrateUploadsToObjectStorage, :sidekiq do
     end
 
     it 'is guarded by .sanity_check!' do
+      expect(BackgroundMigrationWorker).to receive(:perform_async)
       expect(described_class).to receive(:sanity_check!)
 
       enqueue!
@@ -67,4 +66,48 @@ describe Gitlab::BackgroundMigration::MigrateUploadsToObjectStorage, :sidekiq do
       end
     end
   end
+
+  describe '#perform' do
+    def perform
+      described_class.enqueue!(uploads, mounted_as, to_store)
+    end
+
+    # rubocop:disable Style/MultilineIfModifier
+    shared_examples 'outputs correctly' do |success: 0, failures: 0|
+      total = success + failures
+
+      it 'outputs the reports' do
+        expect(Rails.logger).to receive(:info).with(%r{Migrated #{success}/#{total} files})
+
+        perform
+      end if success > 0
+
+      it 'outputs upload failures' do
+        expect(Rails.logger).to receive(:warn).with(/Error .* I am a teapot/)
+
+        perform
+      end if failures > 0
+    end
+    # rubocop:enable Style/MultilineIfModifier
+
+    it_behaves_like 'outputs correctly', success: 10
+
+    it 'migrates files' do
+      perform
+
+      aggregate_failures do
+        projects.each do |project|
+          expect(project.reload.avatar.upload.local?).to be_falsey
+        end
+      end
+    end
+
+    context 'migration is unsuccessful' do
+      before do
+        allow_any_instance_of(ObjectStorage::Concern).to receive(:migrate!).and_raise(CarrierWave::UploadError, "I am a teapot.")
+      end
+
+      it_behaves_like 'outputs correctly', failures: 10
+    end
+  end
 end

spec/tasks/gitlab/uploads_rake_spec.rb

@@ -36,6 +36,8 @@ describe 'gitlab:uploads rake tasks' do
       stub_env('BATCH', batch_size.to_s)
       stub_uploads_object_storage(uploader_class)
       Rake.application.rake_require 'tasks/gitlab/uploads/migrate'
+
+      allow(BackgroundMigrationWorker).to receive(:perform_async)
     end
 
     def run
@@ -43,42 +45,10 @@ describe 'gitlab:uploads rake tasks' do
       run_rake_task("gitlab:uploads:migrate", *args)
     end
 
-    shared_examples 'outputs correctly' do |success: 0, failures: 0|
-      total = success + failures
-
-      it 'outputs the results for each batch' do
-        batch_count = [batch_size, total].min
-
-        expect { run }.to output(%r{Migrated #{batch_count}/#{batch_count} files}).to_stdout
-      end if success > 0 # rubocop:disable Style/MultilineIfModifier
-
-      it 'outputs the results for the task' do
-        expect { run }.to output(%r{Migrated #{success}/#{total} files}).to_stdout
-      end if success > 0 # rubocop:disable Style/MultilineIfModifier
-
-      it 'outputs upload failures' do
-        expect { run }.to output(/Error .* I am a teapot/).to_stdout
-      end if failures > 0 # rubocop:disable Style/MultilineIfModifier
-    end
+    it 'enqueue jobs in batch' do
+      expect(BackgroundMigrationWorker).to receive(:perform_async).exactly(4).times
 
-    it_behaves_like 'outputs correctly', success: 10
-
-    it 'migrates files' do
       run
-
-      aggregate_failures do
-        projects.each do |project|
-          expect(project.reload.avatar.upload.local?).to be_falsey
-        end
-      end
-    end
-
-    context 'migration is unsuccessful' do
-      before do
-        allow_any_instance_of(ObjectStorage::Concern).to receive(:migrate!).and_raise(CarrierWave::UploadError, "I am a teapot.")
-      end
-
-      it_behaves_like 'outputs correctly', failures: 10
     end
   end
 end