Tests for use_replica_if_possible

Issue https://gitlab.com/gitlab-org/gitlab/-/issues/322133

ee/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb

@@ -108,26 +108,53 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
   # We have an extra test for #transaction here to make sure that nested queries
   # are also sent to a primary.
   describe '#transaction' do
-    after do
-      Gitlab::Database::LoadBalancing::Session.clear_session
+    let(:session) { double(:session) }
+
+    before do
+      allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+        .and_return(session)
     end
 
-    it 'runs the transaction and any nested queries on the primary' do
-      primary = double(:connection)
+    context 'session prefers to use a replica' do
+      before do
+        allow(session).to receive(:use_replica?).and_return(true)
+        allow(session).to receive(:use_primary?).and_return(false)
+      end
+
+      it 'runs the transaction and any nested queries on the replica' do
+        replica = double(:connection)
+
+        allow(replica).to receive(:transaction).and_yield
+        allow(replica).to receive(:select)
 
-      allow(primary).to receive(:transaction).and_yield
-      allow(primary).to receive(:select)
+        expect(proxy.load_balancer).to receive(:read)
+          .twice.and_yield(replica)
+        expect(proxy.load_balancer).not_to receive(:read_write)
+        expect(session).not_to receive(:write!)
+
+        proxy.transaction { proxy.select('true') }
+      end
+    end
 
-      expect(proxy.load_balancer).to receive(:read_write)
-        .twice.and_yield(primary)
+    context 'session does not prefer to use a replica' do
+      before do
+        allow(session).to receive(:use_replica?).and_return(false)
+        allow(session).to receive(:use_primary?).and_return(true)
+      end
 
-      # This expectation is put in place to ensure no read is performed.
-      expect(proxy.load_balancer).not_to receive(:read)
+      it 'runs the transaction and any nested queries on the primary and stick to it' do
+        primary = double(:connection)
 
-      proxy.transaction { proxy.select('true') }
+        allow(primary).to receive(:transaction).and_yield
+        allow(primary).to receive(:select)
 
-      expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?)
-        .to eq(true)
+        expect(proxy.load_balancer).to receive(:read_write)
+          .twice.and_yield(primary)
+        expect(proxy.load_balancer).not_to receive(:read)
+        expect(session).to receive(:write!)
+
+        proxy.transaction { proxy.select('true') }
+      end
     end
   end
 
@@ -147,6 +174,32 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
       expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) }
         .not_to raise_error
     end
+
+    context 'current session prefers to use a replica' do
+      let(:session) { double(:session) }
+
+      before do
+        allow(Gitlab::Database::LoadBalancing::Session).to receive(:current)
+          .and_return(session)
+        allow(session).to receive(:use_replica?).and_return(true)
+        allow(session).to receive(:use_primary?).and_return(false)
+      end
+
+      it 'runs the query on the replica' do
+        expect(proxy).to receive(:read_using_load_balancer).with(:foo, ['foo'])
+
+        proxy.foo('foo')
+      end
+
+      it 'properly forwards trailing hash arguments' do
+        allow(proxy.load_balancer).to receive(:read)
+
+        expect(proxy).to receive(:read_using_load_balancer).and_call_original
+
+        expect { proxy.case_sensitive_comparison(:table, :attribute, :column, { value: :value, format: :format }) }
+          .not_to raise_error
+      end
+    end
   end
 
   describe '#read_using_load_balancer' do

ee/spec/lib/gitlab/database/load_balancing/session_spec.rb

@@ -138,4 +138,144 @@ RSpec.describe Gitlab::Database::LoadBalancing::Session do
       expect(instance).to be_using_primary
     end
   end
+
+  describe '#use_replica_if_possible' do
+    let(:instance) { described_class.new }
+
+    it 'uses replica during block' do
+      expect do |blk|
+        instance.use_replica_if_possible do
+          expect(instance.use_replica?).to eq(true)
+
+          # call yield probe
+          blk.to_proc.call
+        end
+      end.to yield_control
+
+      expect(instance.use_replica?).to eq(false)
+    end
+
+    it 'restores state after use' do
+      expect do |blk|
+        instance.use_replica_if_possible do
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(true)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+
+          expect(instance.use_replica?).to eq(true)
+        end
+      end.to yield_control
+
+      expect(instance.use_replica?).to eq(false)
+    end
+
+    context 'when primary was used before' do
+      before do
+        instance.use_primary!
+      end
+
+      it 'uses primary during block' do
+        expect(instance.use_replica?).to eq(false)
+
+        expect do |blk|
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(false)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+        end.to yield_control
+
+        expect(instance.use_replica?).to eq(false)
+      end
+    end
+
+    context 'when a write was performed before' do
+      before do
+        instance.write!
+      end
+
+      it 'uses primary during block' do
+        expect(instance.use_replica?).to eq(false)
+
+        expect do |blk|
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(false)
+
+            # call yield probe
+            blk.to_proc.call
+          end
+        end.to yield_control
+
+        expect(instance.use_replica?).to eq(false)
+      end
+    end
+
+    context 'when primary was used inside the block' do
+      it 'uses primary aterward' do
+        expect(instance.use_replica?).to eq(false)
+
+        instance.use_replica_if_possible do
+          expect(instance.use_replica?).to eq(true)
+
+          instance.use_primary!
+
+          expect(instance.use_replica?).to eq(false)
+        end
+
+        expect(instance.use_replica?).to eq(false)
+      end
+
+      it 'restores state after use' do
+        instance.use_replica_if_possible do
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(true)
+
+            instance.use_primary!
+
+            expect(instance.use_replica?).to eq(false)
+          end
+
+          expect(instance.use_replica?).to eq(false)
+        end
+
+        expect(instance.use_replica?).to eq(false)
+      end
+    end
+
+    context 'when a write was performed inside the block' do
+      it 'uses primary aterward' do
+        expect(instance.use_replica?).to eq(false)
+
+        instance.use_replica_if_possible do
+          expect(instance.use_replica?).to eq(true)
+
+          instance.write!
+
+          expect(instance.use_replica?).to eq(false)
+        end
+
+        expect(instance.use_replica?).to eq(false)
+      end
+
+      it 'restores state after use' do
+        instance.use_replica_if_possible do
+          instance.use_replica_if_possible do
+            expect(instance.use_replica?).to eq(true)
+
+            instance.write!
+
+            expect(instance.use_replica?).to eq(false)
+          end
+
+          expect(instance.use_replica?).to eq(false)
+        end
+
+        expect(instance.use_replica?).to eq(false)
+      end
+    end
+  end
 end