1 files changed, 74 insertions, 24 deletions
diff --git a/spec/lib/gitlab/import_export/lfs_restorer_spec.rb b/spec/lib/gitlab/import_export/lfs_restorer_spec.rb
index 70eeb9ee66b..2b0bdb909ae 100644
--- a/spec/lib/gitlab/import_export/lfs_restorer_spec.rb
+++ b/spec/lib/gitlab/import_export/lfs_restorer_spec.rb
@@ -6,6 +6,7 @@ describe Gitlab::ImportExport::LfsRestorer do
   let(:export_path) { "#{Dir.tmpdir}/lfs_object_restorer_spec" }
   let(:project) { create(:project) }
   let(:shared) { project.import_export_shared }
+  let(:saver) { Gitlab::ImportExport::LfsSaver.new(project: project, shared: shared) }
   subject(:restorer) { described_class.new(project: project, shared: shared) }
 
   before do
@@ -19,49 +20,98 @@ describe Gitlab::ImportExport::LfsRestorer do
 
   describe '#restore' do
     context 'when the archive contains lfs files' do
-      let(:dummy_lfs_file_path) { File.join(shared.export_path, 'lfs-objects', 'dummy') }
-
-      def create_lfs_object_with_content(content)
-        dummy_lfs_file = Tempfile.new('existing')
-        File.write(dummy_lfs_file.path, content)
-        size = dummy_lfs_file.size
-        oid = LfsObject.calculate_oid(dummy_lfs_file.path)
-        LfsObject.create!(oid: oid, size: size, file: dummy_lfs_file)
+      let(:lfs_object) { create(:lfs_object, :correct_oid, :with_file) }
+
+      # Use the LfsSaver to save data to be restored
+      def save_lfs_data
+        %w(project wiki).each do |repository_type|
+          create(
+            :lfs_objects_project,
+            project: project,
+            repository_type: repository_type,
+            lfs_object: lfs_object
+          )
+        end
+
+        saver.save
+
+        project.lfs_objects.delete_all
       end
 
       before do
-        FileUtils.mkdir_p(File.dirname(dummy_lfs_file_path))
-        File.write(dummy_lfs_file_path, 'not very large')
-        allow(restorer).to receive(:lfs_file_paths).and_return([dummy_lfs_file_path])
+        save_lfs_data
+        project.reload
       end
 
-      it 'creates an lfs object for the project' do
-        expect { restorer.restore }.to change { project.reload.lfs_objects.size }.by(1)
+      it 'succeeds' do
+        expect(restorer.restore).to eq(true)
+        expect(shared.errors).to be_empty
       end
 
-      it 'assigns the file correctly' do
+      it 'does not create a new `LfsObject` records, as one already exists' do
+        expect { restorer.restore }.not_to change { LfsObject.count }
+      end
+
+      it 'creates new `LfsObjectsProject` records in order to link the project to the existing `LfsObject`' do
+        expect { restorer.restore }.to change { LfsObjectsProject.count }.by(2)
+      end
+
+      it 'restores the correct `LfsObject` records' do
         restorer.restore
 
-        expect(project.lfs_objects.first.file.read).to eq('not very large')
+        expect(project.lfs_objects).to contain_exactly(lfs_object)
       end
 
-      it 'links an existing LFS object if it existed' do
-        lfs_object = create_lfs_object_with_content('not very large')
+      it 'restores the correct `LfsObjectsProject` records for the project' do
+        restorer.restore
 
+        expect(
+          project.lfs_objects_projects.pluck(:repository_type)
+        ).to contain_exactly('project', 'wiki')
+      end
+
+      it 'assigns the file correctly' do
         restorer.restore
 
-        expect(project.lfs_objects).to include(lfs_object)
+        expect(project.lfs_objects.first.file.read).to eq(lfs_object.file.read)
       end
 
-      it 'succeeds' do
-        expect(restorer.restore).to be_truthy
-        expect(shared.errors).to be_empty
+      context 'when there is not an existing `LfsObject`' do
+        before do
+          lfs_object.destroy
+        end
+
+        it 'creates a new lfs object' do
+          expect { restorer.restore }.to change { LfsObject.count }.by(1)
+        end
+
+        it 'stores the upload' do
+          expect_any_instance_of(LfsObjectUploader).to receive(:store!)
+
+          restorer.restore
+        end
       end
 
-      it 'stores the upload' do
-        expect_any_instance_of(LfsObjectUploader).to receive(:store!)
+      context 'when there is no lfs-objects.json file' do
+        before do
+          json_file = File.join(shared.export_path, ::Gitlab::ImportExport.lfs_objects_filename)
 
-        restorer.restore
+          FileUtils.rm_rf(json_file)
+        end
+
+        it 'restores the correct `LfsObject` records' do
+          restorer.restore
+
+          expect(project.lfs_objects).to contain_exactly(lfs_object)
+        end
+
+        it 'restores a single `LfsObjectsProject` record for the project with "project" for the `repository_type`' do
+          restorer.restore
+
+          expect(
+            project.lfs_objects_projects.pluck(:repository_type)
+          ).to contain_exactly('project')
+        end
       end
     end