Backport EE changes to make test possible when prepending modules

3 files changed, 18 insertions, 6 deletions

diff --git a/app/services/projects/hashed_storage/migrate_attachments_service.rb b/app/services/projects/hashed_storage/migrate_attachments_service.rb
index 93f44110605..8cac6221a96 100644
--- a/app/services/projects/hashed_storage/migrate_attachments_service.rb
+++ b/app/services/projects/hashed_storage/migrate_attachments_service.rb
@@ -17,6 +17,10 @@ module Projects
 
         move_folder!(old_path, new_path)
         project.save!
+
+        if block_given?
+          yield
+        end
       end
 
       private
diff --git a/app/services/projects/hashed_storage/migrate_repository_service.rb b/app/services/projects/hashed_storage/migrate_repository_service.rb
index da43877a185..7212e7524ab 100644
--- a/app/services/projects/hashed_storage/migrate_repository_service.rb
+++ b/app/services/projects/hashed_storage/migrate_repository_service.rb
@@ -35,6 +35,10 @@ module Projects
         project.repository_read_only = false
         project.save!
 
+        if result && block_given?
+          yield
+        end
+
         result
       end
 
diff --git a/spec/services/projects/hashed_storage_migration_service_spec.rb b/spec/services/projects/hashed_storage_migration_service_spec.rb
index 28b6daf217e..466f0b5d7c2 100644
--- a/spec/services/projects/hashed_storage_migration_service_spec.rb
+++ b/spec/services/projects/hashed_storage_migration_service_spec.rb
@@ -6,32 +6,36 @@ describe Projects::HashedStorageMigrationService do
 
   describe '#execute' do
     context 'repository migration' do
+      let(:repository_service) { Projects::HashedStorage::MigrateRepositoryService.new(project, subject.logger) }
+
       it 'delegates migration to Projects::HashedStorage::MigrateRepositoryService' do
-        expect(Projects::HashedStorage::MigrateRepositoryService).to receive(:new).with(project, subject.logger).and_call_original
-        expect_any_instance_of(Projects::HashedStorage::MigrateRepositoryService).to receive(:execute)
+        expect(Projects::HashedStorage::MigrateRepositoryService).to receive(:new).with(project, subject.logger).and_return(repository_service)
+        expect(repository_service).to receive(:execute)
 
         service.execute
       end
 
       it 'does not delegate migration if repository is already migrated' do
         project.storage_version = ::Project::LATEST_STORAGE_VERSION
-        expect_any_instance_of(Projects::HashedStorage::MigrateRepositoryService).not_to receive(:execute)
+        expect(Projects::HashedStorage::MigrateRepositoryService).not_to receive(:new)
 
         service.execute
       end
     end
 
     context 'attachments migration' do
+      let(:attachments_service) { Projects::HashedStorage::MigrateAttachmentsService.new(project, subject.logger) }
+
       it 'delegates migration to Projects::HashedStorage::MigrateRepositoryService' do
-        expect(Projects::HashedStorage::MigrateAttachmentsService).to receive(:new).with(project, subject.logger).and_call_original
-        expect_any_instance_of(Projects::HashedStorage::MigrateAttachmentsService).to receive(:execute)
+        expect(Projects::HashedStorage::MigrateAttachmentsService).to receive(:new).with(project, subject.logger).and_return(attachments_service)
+        expect(attachments_service).to receive(:execute)
 
         service.execute
       end
 
       it 'does not delegate migration if attachments are already migrated' do
         project.storage_version = ::Project::LATEST_STORAGE_VERSION
-        expect_any_instance_of(Projects::HashedStorage::MigrateAttachmentsService).not_to receive(:execute)
+        expect(Projects::HashedStorage::MigrateAttachmentsService).not_to receive(:new)
 
         service.execute
       end