1 files changed, 152 insertions, 0 deletions
diff --git a/spec/models/concerns/atomic_internal_id_spec.rb b/spec/models/concerns/atomic_internal_id_spec.rb
index 5ee3c012dc9..35b0f107676 100644
--- a/spec/models/concerns/atomic_internal_id_spec.rb
+++ b/spec/models/concerns/atomic_internal_id_spec.rb
@@ -87,6 +87,158 @@ RSpec.describe AtomicInternalId do
     end
   end
 
+  describe '#clear_scope_iid!' do
+    context 'when no ensure_if condition is given' do
+      it 'clears automatically set IIDs' do
+        expect(milestone).to receive(:clear_project_iid!).and_call_original
+
+        expect_iid_to_be_set_and_rollback(milestone)
+
+        expect(milestone.iid).to be_nil
+      end
+
+      it 'does not clear manually set IIDS' do
+        milestone.iid = external_iid
+
+        expect(milestone).to receive(:clear_project_iid!).and_call_original
+
+        expect_iid_to_be_set_and_rollback(milestone)
+
+        expect(milestone.iid).to eq(external_iid)
+      end
+    end
+
+    context 'when an ensure_if condition is given' do
+      let(:test_class) do
+        Class.new(ApplicationRecord) do
+          include AtomicInternalId
+          include Importable
+
+          self.table_name = :milestones
+
+          belongs_to :project
+
+          has_internal_id :iid, scope: :project, track_if: -> { !importing }, ensure_if: -> { !importing }
+
+          def self.name
+            'TestClass'
+          end
+        end
+      end
+
+      let(:instance) { test_class.new(milestone.attributes) }
+
+      context 'when the ensure_if condition evaluates to true' do
+        it 'clears automatically set IIDs' do
+          expect(instance).to receive(:clear_project_iid!).and_call_original
+
+          expect_iid_to_be_set_and_rollback(instance)
+
+          expect(instance.iid).to be_nil
+        end
+
+        it 'does not clear manually set IIDs' do
+          instance.iid = external_iid
+
+          expect(instance).to receive(:clear_project_iid!).and_call_original
+
+          expect_iid_to_be_set_and_rollback(instance)
+
+          expect(instance.iid).to eq(external_iid)
+        end
+      end
+
+      context 'when the ensure_if condition evaluates to false' do
+        before do
+          instance.importing = true
+        end
+
+        it 'does not clear IIDs' do
+          instance.iid = external_iid
+
+          expect(instance).not_to receive(:clear_project_iid!)
+
+          expect_iid_to_be_set_and_rollback(instance)
+
+          expect(instance.iid).to eq(external_iid)
+        end
+      end
+    end
+
+    def expect_iid_to_be_set_and_rollback(instance)
+      ActiveRecord::Base.transaction(requires_new: true) do
+        instance.save!
+
+        expect(instance.iid).not_to be_nil
+
+        raise ActiveRecord::Rollback
+      end
+    end
+  end
+
+  describe '#validate_scope_iid_exists!' do
+    let(:test_class) do
+      Class.new(ApplicationRecord) do
+        include AtomicInternalId
+        include Importable
+
+        self.table_name = :milestones
+
+        belongs_to :project
+
+        def self.name
+          'TestClass'
+        end
+      end
+    end
+
+    let(:instance) { test_class.new(milestone.attributes) }
+
+    before do
+      test_class.has_internal_id :iid, scope: :project, presence: presence, ensure_if: -> { !importing }
+
+      instance.importing = true
+    end
+
+    context 'when the presence flag is set' do
+      let(:presence) { true }
+
+      it 'raises an error for blank iids on create' do
+        expect do
+          instance.save!
+        end.to raise_error(described_class::MissingValueError, 'iid was unexpectedly blank!')
+      end
+
+      it 'raises an error for blank iids on update' do
+        instance.iid = 100
+        instance.save!
+
+        instance.iid = nil
+
+        expect do
+          instance.save!
+        end.to raise_error(described_class::MissingValueError, 'iid was unexpectedly blank!')
+      end
+    end
+
+    context 'when the presence flag is not set' do
+      let(:presence) { false }
+
+      it 'does not raise an error for blank iids on create' do
+        expect { instance.save! }.not_to raise_error
+      end
+
+      it 'does not raise an error for blank iids on update' do
+        instance.iid = 100
+        instance.save!
+
+        instance.iid = nil
+
+        expect { instance.save! }.not_to raise_error
+      end
+    end
+  end
+
   describe '.with_project_iid_supply' do
     let(:iid) { 100 }