Move logic for sample data in relation factory

- Fix specs for date calculator
- Fix wrong guard logic
- Fix rubocop offences
- Extract logic to sample relation factory
- Add Sample::RelationFactory
- Rename SampleDataRelationTreeRestorer
- Fix Importer
- Add param for mark_as_consumed

lib/gitlab/import_export/importer.rb

@@ -55,9 +55,17 @@ module Gitlab
       end
 
       def project_tree
-        @project_tree ||= Gitlab::ImportExport::Project::TreeRestorer.new(user: current_user,
-                                                                        shared: shared,
-                                                                        project: project)
+        @project_tree ||= project_tree_class.new(user: current_user,
+                                                 shared: shared,
+                                                 project: project)
+      end
+
+      def project_tree_class
+        sample_data_template? ? Gitlab::ImportExport::Project::Sample::TreeRestorer : Gitlab::ImportExport::Project::TreeRestorer
+      end
+
+      def sample_data_template?
+        project&.import_data&.data&.dig('sample_data')
       end
 
       def avatar_restorer

lib/gitlab/import_export/project/sample/relation_factory.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module ImportExport
+    module Project
+      module Sample
+        class RelationFactory < Project::RelationFactory
+          DATE_MODELS = %i[issues milestones].freeze
+
+          def initialize(date_calculator:, **args)
+            super(**args)
+
+            @date_calculator = date_calculator
+          end
+
+          private
+
+          def setup_models
+            super
+
+            # Override due date attributes in data hash for Sample Data templates
+            # Dates are moved by taking the closest one to average and moving that (and rest around it) to the date of import
+            override_date_attributes
+          end
+
+          def override_date_attributes
+            return unless DATE_MODELS.include?(@relation_name)
+
+            @relation_hash['start_date'] = calculate_by_closest_date(@relation_hash['start_date']&.to_time)
+            @relation_hash['due_date'] = calculate_by_closest_date(@relation_hash['due_date']&.to_time)
+          end
+
+          def calculate_by_closest_date(date)
+            return unless date
+
+            @date_calculator.calculate_by_closest_date_to_average(date)
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/import_export/project/sample/relation_tree_restorer.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module ImportExport
+    module Project
+      module Sample
+        class RelationTreeRestorer < ImportExport::RelationTreeRestorer
+          def initialize(*args)
+            super
+
+            @date_calculator = Gitlab::ImportExport::Project::Sample::DateCalculator.new(dates)
+          end
+
+          private
+
+          def relation_factory_params(*args)
+            super.merge(date_calculator: @date_calculator)
+          end
+
+          def dates
+            return [] if relation_reader.legacy?
+
+            RelationFactory::DATE_MODELS.flat_map do |tag|
+              relation_reader.consume_relation(@importable_path, tag, mark_as_consumed: false).map do |model|
+                model.first['due_date']
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/import_export/project/sample/sample_data_relation_tree_restorer.rb

-# frozen_string_literal: true
-
-module Gitlab
-  module ImportExport
-    module Project
-      module Sample
-        class SampleDataRelationTreeRestorer < RelationTreeRestorer
-          DATE_MODELS = %i[issues milestones].freeze
-
-          def initialize(*args)
-            super
-
-            date_calculator
-          end
-
-          private
-
-          def build_relation(relation_key, relation_definition, data_hash)
-            # Override due date attributes in data hash for Sample Data templates
-            # Dates are moved by taking the closest one to average and moving that (and rest around it) to the date of import
-            # TODO: To move this logic to RelationFactory (see: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/41699#note_430465333)
-            override_date_attributes!(relation_key, data_hash)
-            super
-          end
-
-          def override_date_attributes!(relation_key, data_hash)
-            return unless DATE_MODELS.include?(relation_key.to_sym)
-
-            data_hash['start_date'] = date_calculator.calculate_by_closest_date_to_average(data_hash['start_date'].to_time) unless data_hash['start_date'].nil?
-            data_hash['due_date'] = date_calculator.calculate_by_closest_date_to_average(data_hash['due_date'].to_time) unless data_hash['due_date'].nil?
-          end
-
-          def dates
-            return if relation_reader.legacy?
-
-            DATE_MODELS.flat_map do |tag|
-              relation_reader.consume_relation(@importable_path, tag, mark_as_consumed: false).map do |model|
-                model.first['due_date']
-              end
-            end
-          end
-
-          def date_calculator
-            @date_calculator ||= Gitlab::ImportExport::Project::Sample::DateCalculator.new(dates)
-          end
-        end
-      end
-    end
-  end
-end

lib/gitlab/import_export/project/sample/tree_restorer.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module ImportExport
+    module Project
+      module Sample
+        class TreeRestorer < Project::TreeRestorer
+          def relation_tree_restorer_class
+            RelationTreeRestorer
+          end
+
+          def relation_factory
+            RelationFactory
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/import_export/project/tree_restorer.rb

@@ -85,11 +85,7 @@ module Gitlab
         end
 
         def relation_tree_restorer_class
-          sample_data_template? ? Sample::SampleDataRelationTreeRestorer : RelationTreeRestorer
-        end
-
-        def sample_data_template?
-          @project&.import_data&.data&.dig('sample_data')
+          RelationTreeRestorer
         end
 
         def members_mapper

spec/lib/gitlab/import_export/importer_spec.rb

@@ -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

spec/lib/gitlab/import_export/project/sample/relation_factory_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::ImportExport::Project::Sample::RelationFactory do
+  let(:group) { create(:group) }
+  let(:project) { create(:project, :repository, group: group) }
+  let(:members_mapper) { double('members_mapper').as_null_object }
+  let(:admin) { create(:admin) }
+  let(:importer_user) { admin }
+  let(:excluded_keys) { [] }
+  let(:date_calculator) { instance_double(Gitlab::ImportExport::Project::Sample::DateCalculator) }
+  let(:original_project_id) { 8 }
+  let(:start_date) { Time.current - 30.days }
+  let(:due_date) { Time.current - 20.days }
+  let(:created_object) do
+    described_class.create( # rubocop:disable Rails/SaveBang
+      relation_sym: relation_sym,
+      relation_hash: relation_hash,
+      object_builder: Gitlab::ImportExport::Project::ObjectBuilder,
+      members_mapper: members_mapper,
+      user: importer_user,
+      importable: project,
+      excluded_keys: excluded_keys,
+      date_calculator: date_calculator
+    )
+  end
+
+  context 'issue object' do
+    let(:relation_sym) { :issues }
+    let(:id) { 999 }
+
+    let(:relation_hash) do
+      {
+        'id' => id,
+        'title' => 'Necessitatibus magnam qui at velit consequatur perspiciatis.',
+        'project_id' => original_project_id,
+        'created_at' => '2016-08-12T09:41:03.462Z',
+        'updated_at' => '2016-08-12T09:41:03.462Z',
+        'description' => 'Molestiae corporis magnam et fugit aliquid nulla quia.',
+        'state' => 'closed',
+        'position' => 0,
+        'confidential' => false,
+        'due_date' => due_date
+      }
+    end
+
+    before do
+      allow(date_calculator).to receive(:closest_date_to_average) { Time.current - 10.days }
+      allow(date_calculator).to receive(:calculate_by_closest_date_to_average)
+    end
+
+    it 'correctly updated due date', :aggregate_failures do
+      expect(date_calculator).to receive(:calculate_by_closest_date_to_average)
+        .with(relation_hash['due_date']).and_return(due_date - 10.days)
+
+      expect(created_object.due_date).to eq((due_date - 10.days).to_date)
+    end
+  end
+
+  context 'milestone object' do
+    let(:relation_sym) { :milestones }
+    let(:id) { 1001 }
+
+    let(:relation_hash) do
+      {
+        'id' => id,
+        'title' => 'v3.0',
+        'project_id' => original_project_id,
+        'created_at' => '2016-08-12T09:41:03.462Z',
+        'updated_at' => '2016-08-12T09:41:03.462Z',
+        'description' => 'Rerum at autem exercitationem ea voluptates harum quam placeat.',
+        'state' => 'closed',
+        'start_date' => start_date,
+        'due_date' => due_date
+      }
+    end
+
+    before do
+      allow(date_calculator).to receive(:closest_date_to_average).twice { Time.current - 10.days }
+      allow(date_calculator).to receive(:calculate_by_closest_date_to_average).twice
+    end
+
+    it 'correctly updated due date', :aggregate_failures do
+      expect(date_calculator).to receive(:calculate_by_closest_date_to_average)
+        .with(relation_hash['due_date']).and_return(due_date - 10.days)
+
+      expect(created_object.due_date).to eq((due_date - 10.days).to_date)
+    end
+
+    it 'correctly updated start date', :aggregate_failures do
+      expect(date_calculator).to receive(:calculate_by_closest_date_to_average)
+        .with(relation_hash['start_date']).and_return(start_date - 20.days)
+
+      expect(created_object.start_date).to eq((start_date - 20.days).to_date)
+    end
+  end
+
+  context 'milestone object' do
+    let(:relation_sym) { :milestones }
+    let(:id) { 1001 }
+
+    let(:relation_hash) do
+      {
+        'id' => id,
+        'title' => 'v3.0',
+        'project_id' => original_project_id,
+        'created_at' => '2016-08-12T09:41:03.462Z',
+        'updated_at' => '2016-08-12T09:41:03.462Z',
+        'description' => 'Rerum at autem exercitationem ea voluptates harum quam placeat.',
+        'state' => 'closed',
+        'start_date' => start_date,
+        'due_date' => due_date
+      }
+    end
+
+    before do
+      allow(date_calculator).to receive(:closest_date_to_average).twice { Time.current - 10.days }
+      allow(date_calculator).to receive(:calculate_by_closest_date_to_average).twice
+    end
+
+    it 'correctly updated due date', :aggregate_failures do
+      expect(date_calculator).to receive(:calculate_by_closest_date_to_average)
+        .with(relation_hash['due_date']).and_return(due_date - 10.days)
+
+      expect(created_object.due_date).to eq((due_date - 10.days).to_date)
+    end
+
+    it 'correctly updated start date', :aggregate_failures do
+      expect(date_calculator).to receive(:calculate_by_closest_date_to_average)
+        .with(relation_hash['start_date']).and_return(start_date - 20.days)
+
+      expect(created_object.start_date).to eq((start_date - 20.days).to_date)
+    end
+  end
+
+  context 'hook object' do
+    let(:relation_sym) { :hooks }
+    let(:id) { 999 }
+    let(:service_id) { 99 }
+    let(:token) { 'secret' }
+
+    let(:relation_hash) do
+      {
+        'id' => id,
+        'url' => 'https://example.json',
+        'project_id' => original_project_id,
+        'created_at' => '2016-08-12T09:41:03.462Z',
+        'updated_at' => '2016-08-12T09:41:03.462Z',
+        'service_id' => service_id,
+        'push_events' => true,
+        'issues_events' => false,
+        'confidential_issues_events' => false,
+        'merge_requests_events' => true,
+        'tag_push_events' => false,
+        'note_events' => true,
+        'enable_ssl_verification' => true,
+        'job_events' => false,
+        'wiki_page_events' => true,
+        'token' => token
+      }
+    end
+
+    it 'does not calculate the closest date to average' do
+      expect(date_calculator).not_to receive(:calculate_by_closest_date_to_average)
+    end
+  end
+end

spec/lib/gitlab/import_export/project/sample/sample_data_relation_tree_restorer_spec.rb → spec/lib/gitlab/import_export/project/sample/relation_tree_restorer_spec.rb

@@ -9,7 +9,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::ImportExport::Project::Sample::SampleDataRelationTreeRestorer do
+RSpec.describe Gitlab::ImportExport::Project::Sample::RelationTreeRestorer do
   include_context 'relation tree restorer shared context'
 
   let(:sample_data_relation_tree_restorer) do
@@ -74,13 +74,26 @@ RSpec.describe Gitlab::ImportExport::Project::Sample::SampleDataRelationTreeRest
     let(:importable_name) { 'project' }
     let(:importable_path) { 'project' }
     let(:object_builder) { Gitlab::ImportExport::Project::ObjectBuilder }
-    let(:relation_factory) { Gitlab::ImportExport::Project::RelationFactory }
+    let(:relation_factory) { Gitlab::ImportExport::Project::Sample::RelationFactory }
     let(:reader) { Gitlab::ImportExport::Reader.new(shared: shared) }
+    let(:path) { 'spec/fixtures/lib/gitlab/import_export/sample_data/tree' }
+    let(:relation_reader) { Gitlab::ImportExport::JSON::NdjsonReader.new(path) }
 
-    context 'using ndjson reader' do
-      let(:path) { 'spec/fixtures/lib/gitlab/import_export/sample_data/tree' }
-      let(:relation_reader) { Gitlab::ImportExport::JSON::NdjsonReader.new(path) }
+    it 'initializes relation_factory with date_calculator as parameter' do
+      expect(Gitlab::ImportExport::Project::Sample::RelationFactory).to receive(:create).with(hash_including(:date_calculator)).at_least(:once).times
+
+      subject
+    end
+
+    context 'when relation tree restorer is initialized' do
+      it 'initializes date calculator with due dates' do
+        expect(Gitlab::ImportExport::Project::Sample::DateCalculator).to receive(:new).with(Array)
 
+        sample_data_relation_tree_restorer
+      end
+    end
+
+    context 'using ndjson reader' do
       it_behaves_like 'import project successfully'
     end
   end

spec/lib/gitlab/import_export/project/tree_restorer_spec.rb

@@ -1040,41 +1040,6 @@ RSpec.describe Gitlab::ImportExport::Project::TreeRestorer do
     it_behaves_like 'project tree restorer work properly', :legacy_reader, true
 
     it_behaves_like 'project tree restorer work properly', :ndjson_reader, true
-
-    context 'Sample Data JSON' do
-      let(:user) { create(:user) }
-      let!(:project) { create(:project, :builds_disabled, :issues_disabled, name: 'project', path: 'project') }
-      let(:project_tree_restorer) { described_class.new(user: user, shared: shared, project: project) }
-
-      before do
-        setup_import_export_config('sample_data')
-        setup_reader(:ndjson_reader)
-      end
-
-      context 'with sample_data_template' do
-        before do
-          allow(project).to receive_message_chain(:import_data, :data, :dig).with('sample_data') { true }
-        end
-
-        it 'initialize SampleDataRelationTreeRestorer' do
-          expect_next_instance_of(Gitlab::ImportExport::Project::Sample::SampleDataRelationTreeRestorer) do |restorer|
-            expect(restorer).to receive(:restore).and_return(true)
-          end
-
-          expect(project_tree_restorer.restore).to eq(true)
-        end
-      end
-
-      context 'without sample_data_template' do
-        it 'initialize RelationTreeRestorer' do
-          expect_next_instance_of(Gitlab::ImportExport::RelationTreeRestorer) do |restorer|
-            expect(restorer).to receive(:restore).and_return(true)
-          end
-
-          expect(project_tree_restorer.restore).to eq(true)
-        end
-      end
-    end
   end
 
   context 'disable ndjson import' do