1 files changed, 17 insertions, 9 deletions
diff --git a/spec/support/shared_examples/controllers/rate_limited_endpoint_shared_examples.rb b/spec/support/shared_examples/controllers/rate_limited_endpoint_shared_examples.rb
index 20edca1ee9f..b34038ca0e4 100644
--- a/spec/support/shared_examples/controllers/rate_limited_endpoint_shared_examples.rb
+++ b/spec/support/shared_examples/controllers/rate_limited_endpoint_shared_examples.rb
@@ -5,7 +5,9 @@
 # - current_user
 # - error_message # optional
 
-RSpec.shared_examples 'rate limited endpoint' do |rate_limit_key:|
+RSpec.shared_examples 'rate limited endpoint' do |rate_limit_key:, graphql: false|
+  let(:error_message) { _('This endpoint has been requested too many times. Try again later.') }
+
   context 'when rate limiter enabled', :freeze_time, :clean_gitlab_redis_rate_limiting do
     let(:expected_logger_attributes) do
       {
@@ -25,8 +27,6 @@ RSpec.shared_examples 'rate limited endpoint' do |rate_limit_key:|
       end
     end
 
-    let(:error_message) { _('This endpoint has been requested too many times. Try again later.') }
-
     before do
       allow(Gitlab::ApplicationRateLimiter).to receive(:threshold).with(rate_limit_key).and_return(1)
     end
@@ -37,12 +37,16 @@ RSpec.shared_examples 'rate limited endpoint' do |rate_limit_key:|
       request
       request
 
-      expect(response).to have_gitlab_http_status(:too_many_requests)
+      if graphql
+        expect_graphql_errors_to_include(error_message)
+      else
+        expect(response).to have_gitlab_http_status(:too_many_requests)
 
-      if example.metadata[:type] == :controller
-        expect(response.body).to eq(error_message)
-      else # it is API spec
-        expect(response.body).to eq({ message: { error: error_message } }.to_json)
+        if response.content_type == 'application/json' # it is API spec
+          expect(response.body).to eq({ message: { error: error_message } }.to_json)
+        else
+          expect(response.body).to eq(error_message)
+        end
       end
     end
   end
@@ -57,7 +61,11 @@ RSpec.shared_examples 'rate limited endpoint' do |rate_limit_key:|
 
       request
 
-      expect(response).not_to have_gitlab_http_status(:too_many_requests)
+      if graphql
+        expect_graphql_errors_to_be_empty
+      else
+        expect(response).not_to have_gitlab_http_status(:too_many_requests)
+      end
     end
   end
 end