1 files changed, 12 insertions, 14 deletions
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/rescue_from_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/rescue_from_spec.rb
index c1537c3bd90..22729edb014 100644
--- a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/rescue_from_spec.rb
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/rescue_from_spec.rb
@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
-RSpec.describe ActiveRecord::GitlabPatches::RescueFrom, :without_sqlite3 do
+RSpec.describe ActiveRecord::GitlabPatches::RescueFrom do
   let(:model_with_rescue_from) do
-    Class.new(ActiveRecord::Base) do
-      rescue_from ActiveRecord::ConnectionNotEstablished, with: :handle_exception
+    Class.new(Project) do
+      rescue_from ActiveRecord::StatementInvalid, with: :handle_exception
 
       class << self
         def handle_exception(exception); end
@@ -12,20 +12,18 @@ RSpec.describe ActiveRecord::GitlabPatches::RescueFrom, :without_sqlite3 do
   end
 
   let(:model_without_rescue_from) do
-    Class.new(ActiveRecord::Base)
+    Class.new(Project)
   end
 
-  it 'triggers rescue_from' do
-    stub_const('ModelWithRescueFrom', model_with_rescue_from)
+  context 'for errors from ActiveRelation.load' do
+    it 'triggers rescue_from' do
+      expect(model_with_rescue_from).to receive(:handle_exception)
 
-    expect(model_with_rescue_from).to receive(:handle_exception)
-
-    expect { model_with_rescue_from.all.load }.not_to raise_error
-  end
-
-  it 'does not trigger rescue_from' do
-    stub_const('ModelWithoutRescueFrom', model_without_rescue_from)
+      expect { model_with_rescue_from.where('BADQUERY').load }.not_to raise_error
+    end
 
-    expect { model_without_rescue_from.all.load }.to raise_error(ActiveRecord::ConnectionNotEstablished)
+    it 'does not trigger rescue_from' do
+      expect { model_without_rescue_from.where('BADQUERY').load }.to raise_error(ActiveRecord::StatementInvalid)
+    end
   end
 end