1 files changed, 88 insertions, 0 deletions
diff --git a/activerecord/test/cases/adapters/postgresql/transaction_test.rb b/activerecord/test/cases/adapters/postgresql/transaction_test.rb
index 9b42d0383d..9821b103df 100644
--- a/activerecord/test/cases/adapters/postgresql/transaction_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/transaction_test.rb
@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "cases/helper"
 require "support/connection_helper"
 require "concurrent/atomic/cyclic_barrier"
@@ -12,6 +14,7 @@ module ActiveRecord
 
     setup do
       @abort, Thread.abort_on_exception = Thread.abort_on_exception, false
+      Thread.report_on_exception, @original_report_on_exception = false, Thread.report_on_exception if Thread.respond_to?(:report_on_exception)
 
       @connection = ActiveRecord::Base.connection
 
@@ -29,6 +32,7 @@ module ActiveRecord
       @connection.drop_table "samples", if_exists: true
 
       Thread.abort_on_exception = @abort
+      Thread.report_on_exception = @original_report_on_exception if Thread.respond_to?(:report_on_exception)
     end
 
     test "raises SerializationFailure when a serialization failure occurs" do
@@ -89,6 +93,90 @@ module ActiveRecord
       end
     end
 
+    test "raises LockWaitTimeout when lock wait timeout exceeded" do
+      skip unless ActiveRecord::Base.connection.postgresql_version >= 90300
+      assert_raises(ActiveRecord::LockWaitTimeout) do
+        s = Sample.create!(value: 1)
+        latch1 = Concurrent::CountDownLatch.new
+        latch2 = Concurrent::CountDownLatch.new
+
+        thread = Thread.new do
+          Sample.transaction do
+            Sample.lock.find(s.id)
+            latch1.count_down
+            latch2.wait
+          end
+        end
+
+        begin
+          Sample.transaction do
+            latch1.wait
+            Sample.connection.execute("SET lock_timeout = 1")
+            Sample.lock.find(s.id)
+          end
+        ensure
+          Sample.connection.execute("SET lock_timeout = DEFAULT")
+          latch2.count_down
+          thread.join
+        end
+      end
+    end
+
+    test "raises QueryCanceled when statement timeout exceeded" do
+      assert_raises(ActiveRecord::QueryCanceled) do
+        s = Sample.create!(value: 1)
+        latch1 = Concurrent::CountDownLatch.new
+        latch2 = Concurrent::CountDownLatch.new
+
+        thread = Thread.new do
+          Sample.transaction do
+            Sample.lock.find(s.id)
+            latch1.count_down
+            latch2.wait
+          end
+        end
+
+        begin
+          Sample.transaction do
+            latch1.wait
+            Sample.connection.execute("SET statement_timeout = 1")
+            Sample.lock.find(s.id)
+          end
+        ensure
+          Sample.connection.execute("SET statement_timeout = DEFAULT")
+          latch2.count_down
+          thread.join
+        end
+      end
+    end
+
+    test "raises QueryCanceled when canceling statement due to user request" do
+      assert_raises(ActiveRecord::QueryCanceled) do
+        s = Sample.create!(value: 1)
+        latch = Concurrent::CountDownLatch.new
+
+        thread = Thread.new do
+          Sample.transaction do
+            Sample.lock.find(s.id)
+            latch.count_down
+            sleep(0.5)
+            conn = Sample.connection
+            pid = conn.query_value("SELECT pid FROM pg_stat_activity WHERE query LIKE '% FOR UPDATE'")
+            conn.execute("SELECT pg_cancel_backend(#{pid})")
+          end
+        end
+
+        begin
+          Sample.transaction do
+            latch.wait
+            Sample.lock.find(s.id)
+          end
+        ensure
+          thread.join
+        end
+      end
+    end
+
     private
 
       def with_warning_suppression