1 files changed, 42 insertions, 0 deletions
diff --git a/spec/models/application_record_spec.rb b/spec/models/application_record_spec.rb
index 24de46cb536..85a6717d259 100644
--- a/spec/models/application_record_spec.rb
+++ b/spec/models/application_record_spec.rb
@@ -132,5 +132,47 @@ RSpec.describe ApplicationRecord do
         end.to raise_error(ActiveRecord::QueryCanceled)
       end
     end
+
+    context 'with database load balancing' do
+      let(:session) { double(:session) }
+
+      before do
+        allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)
+        allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries).and_yield
+      end
+
+      it 'yields control' do
+        expect do |blk|
+          described_class.with_fast_read_statement_timeout(&blk)
+        end.to yield_control.once
+      end
+
+      context 'when the query runs faster than configured timeout' do
+        it 'executes the query without error' do
+          result = nil
+
+          expect do
+            described_class.with_fast_read_statement_timeout(100) do
+              result = described_class.connection.exec_query('SELECT 1')
+            end
+          end.not_to raise_error
+
+          expect(result).not_to be_nil
+        end
+      end
+
+      # This query hangs for 10ms and then gets cancelled.  As there is no
+      # other way to test the timeout for sure, 10ms of waiting seems to be
+      # reasonable!
+      context 'when the query runs longer than configured timeout' do
+        it 'cancels the query and raiss an exception' do
+          expect do
+            described_class.with_fast_read_statement_timeout(10) do
+              described_class.connection.exec_query('SELECT pg_sleep(0.1)')
+            end
+          end.to raise_error(ActiveRecord::QueryCanceled)
+        end
+      end
+    end
   end
 end