fix the rake task from being to strict about Group

3 files changed, 116 insertions, 16 deletions

diff --git a/app/workers/object_storage/migrate_uploads_worker.rb b/app/workers/object_storage/migrate_uploads_worker.rb
index 01ed123e6c8..826618aebdb 100644
--- a/app/workers/object_storage/migrate_uploads_worker.rb
+++ b/app/workers/object_storage/migrate_uploads_worker.rb
@@ -138,21 +138,18 @@ module ObjectStorage
 
     include Report
 
-    def self.enqueue!(uploads, mounted_as, to_store)
-      sanity_check!(uploads, mounted_as)
+    def self.enqueue!(uploads, model_class, mounted_as, to_store)
+      sanity_check!(uploads, model_class, mounted_as)
 
-      perform_async(uploads.ids, mounted_as, to_store)
+      perform_async(uploads.ids, model_class.to_s, mounted_as, to_store)
     end
 
     # We need to be sure all the uploads are for the same uploader and model type
     # and that the mount point exists if provided.
     #
-    def self.sanity_check!(uploads, mounted_as)
+    def self.sanity_check!(uploads, model_class, mounted_as)
       upload = uploads.first
-
       uploader_class = upload.uploader.constantize
-      model_class = uploads.first.model_type.constantize
-
       uploader_types = uploads.map(&:uploader).uniq
       model_types = uploads.map(&:model_type).uniq
       model_has_mount = mounted_as.nil? || model_class.uploaders[mounted_as] == uploader_class
@@ -162,7 +159,8 @@ module ObjectStorage
       raise(SanityCheckError, "Mount point #{mounted_as} not found in #{model_class}.") unless model_has_mount
     end
 
-    def perform(ids, mounted_as, to_store)
+    def perform(ids, model_type, mounted_as, to_store)
+      @model_class = model_type.constantize
       @mounted_as = mounted_as&.to_sym
       @to_store = to_store
 
@@ -178,7 +176,7 @@ module ObjectStorage
     end
 
     def sanity_check!(uploads)
-      self.class.sanity_check!(uploads, @mounted_as)
+      self.class.sanity_check!(uploads, @model_class, @mounted_as)
     end
 
     def build_uploaders(uploads)
diff --git a/lib/tasks/gitlab/uploads/migrate.rake b/lib/tasks/gitlab/uploads/migrate.rake
index e7e656a911b..78e18992a8e 100644
--- a/lib/tasks/gitlab/uploads/migrate.rake
+++ b/lib/tasks/gitlab/uploads/migrate.rake
@@ -13,6 +13,7 @@ namespace :gitlab do
 
     def enqueue_batch(batch, index)
       job = ObjectStorage::MigrateUploadsWorker.enqueue!(batch,
+                                                         @model_class,
                                                          @mounted_as,
                                                          @to_store)
       puts "Enqueued job ##{index}: #{job}"
diff --git a/spec/tasks/gitlab/uploads/migrate_rake_spec.rb b/spec/tasks/gitlab/uploads/migrate_rake_spec.rb
index 53b1e45ba54..6fcfae358ec 100644
--- a/spec/tasks/gitlab/uploads/migrate_rake_spec.rb
+++ b/spec/tasks/gitlab/uploads/migrate_rake_spec.rb
@@ -1,10 +1,9 @@
 require 'rake_helper'
 
 describe 'gitlab:uploads:migrate rake tasks' do
-  let!(:projects) { create_list(:project, 10, :with_avatar) }
-  let(:model_class) { Project }
-  let(:uploader_class) { AvatarUploader }
-  let(:mounted_as) { :avatar }
+  let(:model_class) { nil }
+  let(:uploader_class) { nil }
+  let(:mounted_as) { nil }
   let(:batch_size) { 3 }
 
   before do
@@ -30,11 +29,113 @@ describe 'gitlab:uploads:migrate rake tasks' do
     end
   end
 
-  it_behaves_like 'enqueue jobs in batch', batch: 4
+  context "for AvatarUploader" do
+    let(:uploader_class) { AvatarUploader }
+    let(:mounted_as) { :avatar }
+
+    context "for Project" do
+      let(:model_class) { Project }
+      let!(:projects) { create_list(:project, 10, :with_avatar) }
+
+      it_behaves_like 'enqueue jobs in batch', batch: 4
+
+      context 'Upload has store = nil' do
+        before do
+          Upload.where(model: projects).update_all(store: nil)
+        end
+
+        it_behaves_like 'enqueue jobs in batch', batch: 4
+      end
+    end
+
+    context "for Group" do
+      let(:model_class) { Group }
+
+      before do
+        create_list(:group, 10, :with_avatar)
+      end
+
+      it_behaves_like 'enqueue jobs in batch', batch: 4
+    end
+
+    context "for User" do
+      let(:model_class) { User }
+
+      before do
+        create_list(:user, 10, :with_avatar)
+      end
+
+      it_behaves_like 'enqueue jobs in batch', batch: 4
+    end
+  end
+
+  context "for AttachmentUploader" do
+    let(:uploader_class) { AttachmentUploader }
+
+    context "for Note" do
+      let(:model_class) { Note }
+      let(:mounted_as) { :attachment }
+
+      before do
+        create_list(:note, 10, :with_attachment)
+      end
+
+      it_behaves_like 'enqueue jobs in batch', batch: 4
+    end
+
+    context "for Appearance" do
+      let(:model_class) { Appearance }
+      let(:mounted_as) { :logo }
+
+      before do
+        create(:appearance, :with_logos)
+      end
+
+      %i(logo header_logo).each do |mount|
+        it_behaves_like 'enqueue jobs in batch', batch: 1 do
+          let(:mounted_as) { mount }
+        end
+      end
+    end
+  end
+
+  context "for FileUploader" do
+    let(:uploader_class) { FileUploader }
+    let(:model_class) { Project }
+
+    before do
+      create_list(:project, 10) do |model|
+        uploader_class.new(model)
+          .store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
+      end
+    end
+
+    it_behaves_like 'enqueue jobs in batch', batch: 4
+  end
+
+  context "for PersonalFileUploader" do
+    let(:uploader_class) { PersonalFileUploader }
+    let(:model_class) { PersonalSnippet }
+
+    before do
+      create_list(:personal_snippet, 10) do |model|
+        uploader_class.new(model)
+          .store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
+      end
+    end
+
+    it_behaves_like 'enqueue jobs in batch', batch: 4
+  end
+
+  context "for NamespaceFileUploader" do
+    let(:uploader_class) { NamespaceFileUploader }
+    let(:model_class) { Snippet }
 
-  context 'Upload has store = nil' do
     before do
-      Upload.where(model: projects).update_all(store: nil)
+      create_list(:snippet, 10) do |model|
+        uploader_class.new(model)
+          .store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
+      end
     end
 
     it_behaves_like 'enqueue jobs in batch', batch: 4