1 files changed, 69 insertions, 8 deletions

diff --git a/spec/models/application_record_spec.rb b/spec/models/application_record_spec.rb
index 85a6717d259..f9a05c720a3 100644
--- a/spec/models/application_record_spec.rb
+++ b/spec/models/application_record_spec.rb
@@ -39,13 +39,14 @@ RSpec.describe ApplicationRecord do
 
     let(:suggestion_attributes) { attributes_for(:suggestion).merge!(note_id: note.id) }
 
-    describe '.safe_find_or_create_by' do
+    shared_examples '.safe_find_or_create_by' do
       it 'creates the suggestion avoiding race conditions' do
-        expect(Suggestion).to receive(:find_or_create_by).and_raise(ActiveRecord::RecordNotUnique)
-        allow(Suggestion).to receive(:find_or_create_by).and_call_original
+        existing_suggestion = double(:Suggestion)
 
-        expect { Suggestion.safe_find_or_create_by(suggestion_attributes) }
-          .to change { Suggestion.count }.by(1)
+        expect(Suggestion).to receive(:find_by).and_return(nil, existing_suggestion)
+        expect(Suggestion).to receive(:create).and_raise(ActiveRecord::RecordNotUnique)
+
+        expect(Suggestion.safe_find_or_create_by(suggestion_attributes)).to eq(existing_suggestion)
       end
 
       it 'passes a block to find_or_create_by' do
@@ -62,10 +63,8 @@ RSpec.describe ApplicationRecord do
       end
     end
 
-    describe '.safe_find_or_create_by!' do
+    shared_examples '.safe_find_or_create_by!' do
       it 'creates a record using safe_find_or_create_by' do
-        expect(Suggestion).to receive(:find_or_create_by).and_call_original
-
         expect(Suggestion.safe_find_or_create_by!(suggestion_attributes))
           .to be_a(Suggestion)
       end
@@ -89,6 +88,24 @@ RSpec.describe ApplicationRecord do
           .to raise_error(ActiveRecord::RecordNotFound)
       end
     end
+
+    context 'when optimized_safe_find_or_create_by is enabled' do
+      before do
+        stub_feature_flags(optimized_safe_find_or_create_by: true)
+      end
+
+      it_behaves_like '.safe_find_or_create_by'
+      it_behaves_like '.safe_find_or_create_by!'
+    end
+
+    context 'when optimized_safe_find_or_create_by is disabled' do
+      before do
+        stub_feature_flags(optimized_safe_find_or_create_by: false)
+      end
+
+      it_behaves_like '.safe_find_or_create_by'
+      it_behaves_like '.safe_find_or_create_by!'
+    end
   end
 
   describe '.underscore' do
@@ -105,6 +122,50 @@ RSpec.describe ApplicationRecord do
     end
   end
 
+  describe '.transaction', :delete do
+    it 'opens a new transaction' do
+      expect(described_class.connection.transaction_open?).to be false
+
+      Project.transaction do
+        expect(Project.connection.transaction_open?).to be true
+
+        Project.transaction(requires_new: true) do
+          expect(Project.connection.transaction_open?).to be true
+        end
+      end
+    end
+
+    it 'does not increment a counter when a transaction is not nested' do
+      expect(described_class.connection.transaction_open?).to be false
+
+      expect(::Gitlab::Database::Metrics)
+        .not_to receive(:subtransactions_increment)
+
+      Project.transaction do
+        expect(Project.connection.transaction_open?).to be true
+      end
+
+      Project.transaction(requires_new: true) do
+        expect(Project.connection.transaction_open?).to be true
+      end
+    end
+
+    it 'increments a counter when a nested transaction is created' do
+      expect(described_class.connection.transaction_open?).to be false
+
+      expect(::Gitlab::Database::Metrics)
+        .to receive(:subtransactions_increment)
+        .with('Project')
+        .once
+
+      Project.transaction do
+        Project.transaction(requires_new: true) do
+          expect(Project.connection.transaction_open?).to be true
+        end
+      end
+    end
+  end
+
   describe '.with_fast_read_statement_timeout' do
     context 'when the query runs faster than configured timeout' do
       it 'executes the query without error' do