1 files changed, 29 insertions, 0 deletions
diff --git a/spec/models/application_record_spec.rb b/spec/models/application_record_spec.rb
index 107f707ccd9..7e6ac351e68 100644
--- a/spec/models/application_record_spec.rb
+++ b/spec/models/application_record_spec.rb
@@ -100,4 +100,33 @@ RSpec.describe ApplicationRecord do
       expect(User.where_exists(User.limit(1))).to eq([user])
     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
+        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 raises 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