2 files changed, 129 insertions, 5 deletions
diff --git a/spec/lib/gitlab/auth_spec.rb b/spec/lib/gitlab/auth_spec.rb
index 4f4a27e4c41..af1db2c3455 100644
--- a/spec/lib/gitlab/auth_spec.rb
+++ b/spec/lib/gitlab/auth_spec.rb
@@ -16,20 +16,20 @@ describe Gitlab::Auth do
       expect(subject::DEFAULT_SCOPES).to eq [:api]
     end
 
-    it 'OPTIONAL_SCOPES contains all non-default scopes' do
+    it 'optional_scopes contains all non-default scopes' do
       stub_container_registry_config(enabled: true)
 
-      expect(subject::OPTIONAL_SCOPES).to eq %i[read_user read_registry openid]
+      expect(subject.optional_scopes).to eq %i[read_user read_registry openid]
     end
 
-    context 'REGISTRY_SCOPES' do
+    context 'registry_scopes' do
       context 'when registry is disabled' do
         before do
           stub_container_registry_config(enabled: false)
         end
 
         it 'is empty' do
-          expect(subject::REGISTRY_SCOPES).to eq []
+          expect(subject.registry_scopes).to eq []
         end
       end
 
@@ -39,7 +39,7 @@ describe Gitlab::Auth do
         end
 
         it 'contains all registry related scopes' do
-          expect(subject::REGISTRY_SCOPES).to eq %i[read_registry]
+          expect(subject.registry_scopes).to eq %i[read_registry]
         end
       end
     end
diff --git a/spec/lib/gitlab/gitaly_client_spec.rb b/spec/lib/gitlab/gitaly_client_spec.rb
index a9b861fcff2..9a84d6e6a67 100644
--- a/spec/lib/gitlab/gitaly_client_spec.rb
+++ b/spec/lib/gitlab/gitaly_client_spec.rb
@@ -38,6 +38,130 @@ describe Gitlab::GitalyClient, skip_gitaly_mock: true do
     end
   end
 
+  describe 'allow_n_plus_1_calls' do
+    context 'when RequestStore is enabled', :request_store do
+      it 'returns the result of the allow_n_plus_1_calls block' do
+        expect(described_class.allow_n_plus_1_calls { "result" }).to eq("result")
+      end
+    end
+
+    context 'when RequestStore is not active' do
+      it 'returns the result of the allow_n_plus_1_calls block' do
+        expect(described_class.allow_n_plus_1_calls { "something" }).to eq("something")
+      end
+    end
+  end
+
+  describe 'enforce_gitaly_request_limits?' do
+    def call_gitaly(count = 1)
+      (1..count).each do
+        described_class.enforce_gitaly_request_limits(:test)
+      end
+    end
+
+    context 'when RequestStore is enabled', :request_store do
+      it 'allows up the maximum number of allowed calls' do
+        expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS) }.not_to raise_error
+      end
+
+      context 'when the maximum number of calls has been reached' do
+        before do
+          call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS)
+        end
+
+        it 'fails on the next call' do
+          expect { call_gitaly(1) }.to raise_error(Gitlab::GitalyClient::TooManyInvocationsError)
+        end
+      end
+
+      it 'allows the maximum number of calls to be exceeded within an allow_n_plus_1_calls block' do
+        expect do
+          described_class.allow_n_plus_1_calls do
+            call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1)
+          end
+        end.not_to raise_error
+      end
+
+      context 'when the maximum number of calls has been reached within an allow_n_plus_1_calls block' do
+        before do
+          described_class.allow_n_plus_1_calls do
+            call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS)
+          end
+        end
+
+        it 'allows up to the maximum number of calls outside of an allow_n_plus_1_calls block' do
+          expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS) }.not_to raise_error
+        end
+
+        it 'does not allow the maximum number of calls to be exceeded outside of an allow_n_plus_1_calls block' do
+          expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1) }.to raise_error(Gitlab::GitalyClient::TooManyInvocationsError)
+        end
+      end
+    end
+
+    context 'when RequestStore is not active' do
+      it 'does not raise errors when the maximum number of allowed calls is exceeded' do
+        expect { call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 2) }.not_to raise_error
+      end
+
+      it 'does not fail when the maximum number of calls is exceeded within an allow_n_plus_1_calls block' do
+        expect do
+          described_class.allow_n_plus_1_calls do
+            call_gitaly(Gitlab::GitalyClient::MAXIMUM_GITALY_CALLS + 1)
+          end
+        end.not_to raise_error
+      end
+    end
+  end
+
+  describe 'get_request_count' do
+    context 'when RequestStore is enabled', :request_store do
+      context 'when enforce_gitaly_request_limits is called outside of allow_n_plus_1_calls blocks' do
+        before do
+          described_class.enforce_gitaly_request_limits(:call)
+        end
+
+        it 'counts gitaly calls' do
+          expect(described_class.get_request_count).to eq(1)
+        end
+      end
+
+      context 'when enforce_gitaly_request_limits is called inside and outside of allow_n_plus_1_calls blocks' do
+        before do
+          described_class.enforce_gitaly_request_limits(:call)
+          described_class.allow_n_plus_1_calls do
+            described_class.enforce_gitaly_request_limits(:call)
+          end
+        end
+
+        it 'counts gitaly calls' do
+          expect(described_class.get_request_count).to eq(2)
+        end
+      end
+
+      context 'when reset_counts is called' do
+        before do
+          described_class.enforce_gitaly_request_limits(:call)
+          described_class.reset_counts
+        end
+
+        it 'resets counts' do
+          expect(described_class.get_request_count).to eq(0)
+        end
+      end
+    end
+
+    context 'when RequestStore is not active' do
+      before do
+        described_class.enforce_gitaly_request_limits(:call)
+      end
+
+      it 'returns zero' do
+        expect(described_class.get_request_count).to eq(0)
+      end
+    end
+  end
+
   describe 'feature_enabled?' do
     let(:feature_name) { 'my_feature' }
     let(:real_feature_name) { "gitaly_#{feature_name}" }