1 files changed, 47 insertions, 3 deletions

diff --git a/spec/lib/gitlab/import_export/project/tree_saver_spec.rb b/spec/lib/gitlab/import_export/project/tree_saver_spec.rb
index 374d688576e..f68ec21039d 100644
--- a/spec/lib/gitlab/import_export/project/tree_saver_spec.rb
+++ b/spec/lib/gitlab/import_export/project/tree_saver_spec.rb
@@ -5,6 +5,9 @@ require 'spec_helper'
 RSpec.describe Gitlab::ImportExport::Project::TreeSaver do
   let_it_be(:export_path) { "#{Dir.tmpdir}/project_tree_saver_spec" }
   let_it_be(:exportable_path) { 'project' }
+  let_it_be(:user) { create(:user) }
+  let_it_be(:group) { create(:group) }
+  let_it_be(:project) { setup_project }
 
   shared_examples 'saves project tree successfully' do |ndjson_enabled|
     include ImportExport::CommonUtil
@@ -12,9 +15,6 @@ RSpec.describe Gitlab::ImportExport::Project::TreeSaver do
     subject { get_json(full_path, exportable_path, relation_name, ndjson_enabled) }
 
     describe 'saves project tree attributes' do
-      let_it_be(:user) { create(:user) }
-      let_it_be(:group) { create(:group) }
-      let_it_be(:project) { setup_project }
       let_it_be(:shared) { project.import_export_shared }
 
       let(:relation_name) { :projects }
@@ -402,6 +402,50 @@ RSpec.describe Gitlab::ImportExport::Project::TreeSaver do
     it_behaves_like "saves project tree successfully", true
   end
 
+  context 'when streaming has to retry', :aggregate_failures do
+    let(:shared) { double('shared', export_path: exportable_path) }
+    let(:logger) { Gitlab::Import::Logger.build }
+    let(:serializer) { double('serializer') }
+    let(:error_class) { Net::OpenTimeout }
+    let(:info_params) do
+      {
+        'error.class': error_class,
+        project_name: project.name,
+        project_id: project.id
+      }
+    end
+
+    before do
+      allow(Gitlab::ImportExport::Json::StreamingSerializer).to receive(:new).and_return(serializer)
+    end
+
+    subject(:project_tree_saver) do
+      described_class.new(project: project, current_user: user, shared: shared, logger: logger)
+    end
+
+    it 'retries and succeeds' do
+      call_count = 0
+      allow(serializer).to receive(:execute) do
+        call_count += 1
+        call_count > 1 ? true : raise(error_class, 'execution expired')
+      end
+
+      expect(logger).to receive(:info).with(hash_including(info_params)).once
+
+      expect(project_tree_saver.save).to be(true)
+    end
+
+    it 'retries and does not succeed' do
+      retry_count = 3
+      allow(serializer).to receive(:execute).and_raise(error_class, 'execution expired')
+
+      expect(logger).to receive(:info).with(hash_including(info_params)).exactly(retry_count).times
+      expect(shared).to receive(:error).with(instance_of(error_class))
+
+      expect(project_tree_saver.save).to be(false)
+    end
+  end
+
   def setup_project
     release = create(:release)