1 files changed, 29 insertions, 2 deletions
diff --git a/activerecord/test/cases/adapters/postgresql/transaction_test.rb b/activerecord/test/cases/adapters/postgresql/transaction_test.rb
index ebba0f0c1c..c24dfeb345 100644
--- a/activerecord/test/cases/adapters/postgresql/transaction_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/transaction_test.rb
@@ -120,8 +120,8 @@ module ActiveRecord
       end
     end
 
-    test "raises StatementTimeout when statement timeout exceeded" do
-      assert_raises(ActiveRecord::StatementTimeout) do
+    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
@@ -148,6 +148,33 @@ module ActiveRecord
       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