1 files changed, 16 insertions, 4 deletions
diff --git a/spec/lib/gitlab/import_export/importer_spec.rb b/spec/lib/gitlab/import_export/importer_spec.rb
index dc44296321c..0db038785d3 100644
--- a/spec/lib/gitlab/import_export/importer_spec.rb
+++ b/spec/lib/gitlab/import_export/importer_spec.rb
@@ -12,7 +12,7 @@ RSpec.describe Gitlab::ImportExport::Importer do
   subject(:importer) { described_class.new(project) }
 
   before do
-    allow_any_instance_of(Gitlab::ImportExport).to receive(:storage_path).and_return(test_path)
+    allow(Gitlab::ImportExport).to receive(:storage_path).and_return(test_path)
     allow_any_instance_of(Gitlab::ImportExport::FileImporter).to receive(:remove_import_file)
     stub_uploads_object_storage(FileUploader)
 
@@ -65,10 +65,22 @@ RSpec.describe Gitlab::ImportExport::Importer do
         end
       end
 
-      it 'restores the ProjectTree' do
-        expect(Gitlab::ImportExport::Project::TreeRestorer).to receive(:new).and_call_original
+      context 'with sample_data_template' do
+        it 'initializes the Sample::TreeRestorer' do
+          project.create_or_update_import_data(data: { sample_data: true })
 
-        importer.execute
+          expect(Gitlab::ImportExport::Project::Sample::TreeRestorer).to receive(:new).and_call_original
+
+          importer.execute
+        end
+      end
+
+      context 'without sample_data_template' do
+        it 'initializes the ProjectTree' do
+          expect(Gitlab::ImportExport::Project::TreeRestorer).to receive(:new).and_call_original
+
+          importer.execute
+        end
       end
 
       it 'removes the import file' do