1 files changed, 61 insertions, 22 deletions
diff --git a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
index bf4e3756e0e..53445d73756 100644
--- a/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
@@ -46,41 +46,68 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
 
   describe '.all_caught_up?' do
     let(:lb) { double(:lb) }
+    let(:last_write_location) { 'foo' }
 
     before do
       allow(described_class).to receive(:load_balancer).and_return(lb)
-    end
 
-    it 'returns true if no write location could be found' do
       allow(described_class).to receive(:last_write_location_for)
         .with(:user, 42)
-        .and_return(nil)
+        .and_return(last_write_location)
+    end
+
+    context 'when no write location could be found' do
+      let(:last_write_location) { nil }
 
-      expect(lb).not_to receive(:all_caught_up?)
+      it 'returns true' do
+        allow(described_class).to receive(:last_write_location_for)
+          .with(:user, 42)
+          .and_return(nil)
+
+        expect(lb).not_to receive(:select_up_to_date_host)
 
-      expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+        expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+      end
     end
 
-    it 'returns true, and unsticks if all secondaries have caught up' do
-      allow(described_class).to receive(:last_write_location_for)
-        .with(:user, 42)
-        .and_return('foo')
+    context 'when all secondaries have caught up' do
+      before do
+        allow(lb).to receive(:select_up_to_date_host).with('foo').and_return(true)
+      end
 
-      allow(lb).to receive(:all_caught_up?).with('foo').and_return(true)
+      it 'returns true, and unsticks' do
+        expect(described_class).to receive(:unstick).with(:user, 42)
 
-      expect(described_class).to receive(:unstick).with(:user, 42)
+        expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+      end
+
+      it 'notifies with the proper event payload' do
+        expect(ActiveSupport::Notifications)
+          .to receive(:instrument)
+          .with('caught_up_replica_pick.load_balancing', { result: true })
+          .and_call_original
 
-      expect(described_class.all_caught_up?(:user, 42)).to eq(true)
+        described_class.all_caught_up?(:user, 42)
+      end
     end
 
-    it 'return false if the secondaries have not yet caught up' do
-      allow(described_class).to receive(:last_write_location_for)
-        .with(:user, 42)
-        .and_return('foo')
+    context 'when the secondaries have not yet caught up' do
+      before do
+        allow(lb).to receive(:select_up_to_date_host).with('foo').and_return(false)
+      end
+
+      it 'returns false' do
+        expect(described_class.all_caught_up?(:user, 42)).to eq(false)
+      end
 
-      allow(lb).to receive(:all_caught_up?).with('foo').and_return(false)
+      it 'notifies with the proper event payload' do
+        expect(ActiveSupport::Notifications)
+          .to receive(:instrument)
+          .with('caught_up_replica_pick.load_balancing', { result: false })
+          .and_call_original
 
-      expect(described_class.all_caught_up?(:user, 42)).to eq(false)
+        described_class.all_caught_up?(:user, 42)
+      end
     end
   end
 
@@ -96,7 +123,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
         .with(:user, 42)
         .and_return(nil)
 
-      expect(lb).not_to receive(:all_caught_up?)
+      expect(lb).not_to receive(:select_up_to_date_host)
 
       described_class.unstick_or_continue_sticking(:user, 42)
     end
@@ -106,7 +133,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
         .with(:user, 42)
         .and_return('foo')
 
-      allow(lb).to receive(:all_caught_up?).with('foo').and_return(true)
+      allow(lb).to receive(:select_up_to_date_host).with('foo').and_return(true)
 
       expect(described_class).to receive(:unstick).with(:user, 42)
 
@@ -118,7 +145,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
         .with(:user, 42)
         .and_return('foo')
 
-      allow(lb).to receive(:all_caught_up?).with('foo').and_return(false)
+      allow(lb).to receive(:select_up_to_date_host).with('foo').and_return(false)
 
       expect(Gitlab::Database::LoadBalancing::Session.current)
         .to receive(:use_primary!)
@@ -298,10 +325,22 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
       end
 
       it 'returns true, selects hosts, and unsticks if any secondary has caught up' do
-        expect(lb).to receive(:select_caught_up_hosts).and_return(true)
+        expect(lb).to receive(:select_up_to_date_host).and_return(true)
         expect(described_class).to receive(:unstick).with(:project, 42)
         expect(described_class.select_caught_up_replicas(:project, 42)).to be true
       end
+
+      context 'when :load_balancing_refine_load_balancer_methods FF is disabled' do
+        before do
+          stub_feature_flags(load_balancing_refine_load_balancer_methods: false)
+        end
+
+        it 'returns true, selects hosts, and unsticks if any secondary has caught up' do
+          expect(lb).to receive(:select_caught_up_hosts).and_return(true)
+          expect(described_class).to receive(:unstick).with(:project, 42)
+          expect(described_class.select_caught_up_replicas(:project, 42)).to be true
+        end
+      end
     end
   end
 end