1 files changed, 4 insertions, 2 deletions
diff --git a/spec/support/gitaly.rb b/spec/support/gitaly.rb
index c7e8a39a617..9cf541372b5 100644
--- a/spec/support/gitaly.rb
+++ b/spec/support/gitaly.rb
@@ -1,11 +1,13 @@
 RSpec.configure do |config|
   config.before(:each) do |example|
     if example.metadata[:disable_gitaly]
-      allow(Gitlab::GitalyClient).to receive(:feature_enabled?).and_return(false)
+      # Use 'and_wrap_original' to make sure the arguments are valid
+      allow(Gitlab::GitalyClient).to receive(:feature_enabled?).and_wrap_original { |m, *args| m.call(*args) && false }
     else
       next if example.metadata[:skip_gitaly_mock]
 
-      allow(Gitlab::GitalyClient).to receive(:feature_enabled?).and_return(true)
+      # Use 'and_wrap_original' to make sure the arguments are valid
+      allow(Gitlab::GitalyClient).to receive(:feature_enabled?).and_wrap_original { |m, *args| m.call(*args) || true }
     end
   end
 end