2 files changed, 114 insertions, 1 deletions
diff --git a/spec/support_specs/helpers/stub_method_calls_spec.rb b/spec/support_specs/helpers/stub_method_calls_spec.rb
new file mode 100644
index 00000000000..837a2162bcd
--- /dev/null
+++ b/spec/support_specs/helpers/stub_method_calls_spec.rb
@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe StubMethodCalls do
+  include described_class
+
+  let(:object) do
+    Class.new do
+      def self.test_method
+        'test'
+      end
+
+      def self.test_method_two(response: nil)
+        response || 'test_two'
+      end
+    end
+  end
+
+  describe '#stub_method' do
+    let(:method_to_stub) { :test_method }
+
+    it 'stubs the method response' do
+      stub_method(object, method_to_stub) { true }
+
+      expect(object.send(method_to_stub)).to eq(true)
+    end
+
+    context 'when calling it on an already stubbed method' do
+      before do
+        stub_method(object, method_to_stub) { false }
+      end
+
+      it 'stubs correctly' do
+        stub_method(object, method_to_stub) { true }
+
+        expect(object.send(method_to_stub)).to eq(true)
+      end
+    end
+
+    context 'methods that accept arguments' do
+      it 'stubs correctly' do
+        stub_method(object, method_to_stub) { |a, b| a + b }
+
+        expect(object.send(method_to_stub, 1, 2)).to eq(3)
+      end
+
+      context 'methods that use named arguments' do
+        let(:method_to_stub) { :test_method_two }
+
+        it 'stubs correctly' do
+          stub_method(object, method_to_stub) { |a: 'test'| a }
+
+          expect(object.send(method_to_stub, a: 'testing')).to eq('testing')
+          expect(object.send(method_to_stub)).to eq('test')
+        end
+
+        context 'stubbing non-existent method' do
+          let(:method_to_stub) { :another_method }
+
+          it 'stubs correctly' do
+            stub_method(object, method_to_stub) { |a: 'test'| a }
+
+            expect(object.send(method_to_stub, a: 'testing')).to eq('testing')
+            expect(object.send(method_to_stub)).to eq('test')
+          end
+        end
+      end
+    end
+  end
+
+  describe '#restore_original_method' do
+    before do
+      stub_method(object, :test_method) { true }
+    end
+
+    it 'restores original behaviour' do
+      expect(object.test_method).to eq(true)
+
+      restore_original_method(object, :test_method)
+
+      expect(object.test_method).to eq('test')
+    end
+
+    context 'method is not stubbed' do
+      specify do
+        expect do
+          restore_original_method(object, 'some_other_method')
+        end.to raise_error(NotImplementedError, "some_other_method has not been stubbed on #{object}")
+      end
+    end
+  end
+
+  describe '#restore_original_methods' do
+    before do
+      stub_method(object, :test_method) { true }
+      stub_method(object, :test_method_two) { true }
+    end
+
+    it 'restores original behaviour' do
+      restore_original_methods(object)
+
+      expect(object.test_method).to eq('test')
+      expect(object.test_method_two).to eq('test_two')
+    end
+  end
+end
diff --git a/spec/support_specs/matchers/exceed_query_limit_helpers_spec.rb b/spec/support_specs/matchers/exceed_query_limit_helpers_spec.rb
index 67d87fe3c2f..a6f5b3862a2 100644
--- a/spec/support_specs/matchers/exceed_query_limit_helpers_spec.rb
+++ b/spec/support_specs/matchers/exceed_query_limit_helpers_spec.rb
@@ -56,6 +56,7 @@ RSpec.describe ExceedQueryLimitHelpers do
         TestQueries.where(version: 'x').update_all(version: 'y')
         TestQueries.where(version: 'foobar').count
         TestQueries.where(version: 'z').delete_all
+        Project.where(id: 1).pluck(:title)
       end
     end
 
@@ -71,10 +72,11 @@ RSpec.describe ExceedQueryLimitHelpers do
         TestQueries.count
         TestQueries.where(version: 'y').update_all(version: 'z')
         TestQueries.where(version: 'z').delete_all
+        Project.where(id: 2).pluck(:title)
       end
     end
 
-    it 'merges two query counts' do
+    it 'merges two query counts, showing only diffs' do
       test_matcher = TestMatcher.new
 
       diff = test_matcher.diff_query_counts(
@@ -131,6 +133,10 @@ RSpec.describe ExceedQueryLimitHelpers do
         },
         "RELEASE SAVEPOINT active_record_1" => {
           "" => [0, 1]
+        },
+        "SELECT \"projects\".\"name\" FROM \"projects\"" => {
+          "WHERE \"projects\".\"id\" = 1" => [1, 0],
+          "WHERE \"projects\".\"id\" = 2" => [0, 1]
         }
       })
     end