1 files changed, 124 insertions, 0 deletions
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}" }