Fix logic for consume_relation in reader

- add param for mark_as_consumed

lib/gitlab/import_export/json/ndjson_reader.rb

@@ -29,9 +29,9 @@ module Gitlab
           json_decode(data)
         end
 
-        def consume_relation(importable_path, key)
+        def consume_relation(importable_path, key, mark_as_consumed: true)
           Enumerator.new do |documents|
-            next unless @consumed_relations.add?("#{importable_path}/#{key}")
+            next if mark_as_consumed && !@consumed_relations.add?("#{importable_path}/#{key}")
 
             # This reads from `tree/project/merge_requests.ndjson`
             path = file_path(importable_path, "#{key}.ndjson")
@@ -44,11 +44,6 @@ module Gitlab
           end
         end
 
-        # TODO: Move clear logic into main comsume_relation method (see https://gitlab.com/gitlab-org/gitlab/-/merge_requests/41699#note_430465330)
-        def clear_consumed_relations
-          @consumed_relations.clear
-        end
-
         private
 
         def json_decode(string)

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

@@ -9,7 +9,6 @@ module Gitlab
 
           def initialize(dates)
             @dates = dates.dup
-            @dates.flatten!
             @dates.compact!
             @dates.sort!
             @dates.map! { |date| date.to_time.to_f }

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

@@ -30,13 +30,12 @@ module Gitlab
             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
 
-          # TODO: Move clear logic into main comsume_relation method (see https://gitlab.com/gitlab-org/gitlab/-/merge_requests/41699#note_430465330)
           def dates
-            unless relation_reader.legacy?
-              DATE_MODELS.map do |tag|
-                relation_reader.consume_relation(@importable_path, tag).map { |model| model.first['due_date'] }.tap do
-                  relation_reader.clear_consumed_relations
-                end
+            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

spec/lib/gitlab/import_export/json/ndjson_reader_spec.rb

@@ -67,6 +67,14 @@ RSpec.describe Gitlab::ImportExport::JSON::NdjsonReader do
       it 'yields nothing to the Enumerator' do
         expect(subject.to_a).to eq([])
       end
+
+      context 'with mark_as_consumed: false' do
+        subject { ndjson_reader.consume_relation(importable_path, key, mark_as_consumed: false) }
+
+        it 'yields every relation value to the Enumerator' do
+          expect(subject.count).to eq(1)
+        end
+      end
     end
 
     context 'key has not been consumed' do
@@ -102,14 +110,4 @@ RSpec.describe Gitlab::ImportExport::JSON::NdjsonReader do
       end
     end
   end
-
-  describe '#clear_consumed_relations' do
-    let(:dir_path) { fixture }
-
-    subject { ndjson_reader.clear_consumed_relations }
-
-    it 'returns empty set' do
-      expect(subject).to be_empty
-    end
-  end
 end

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

@@ -13,7 +13,7 @@ RSpec.describe Gitlab::ImportExport::Project::Sample::DateCalculator do
     end
 
     context 'when dates are not empty' do
-      let(:dates) { [[nil, '2020-01-01 00:00:00 +0000'], [nil, '2021-01-01 00:00:00 +0000'], [nil, '2022-01-01 23:59:59 +0000']] }
+      let(:dates) { [nil, '2020-01-01 00:00:00 +0000', '2021-01-01 00:00:00 +0000', nil, '2022-01-01 23:59:59 +0000'] }
 
       it { is_expected.to eq(Time.zone.parse('2021-01-01 00:00:00 +0000')) }
     end