1 files changed, 80 insertions, 0 deletions
diff --git a/spec/lib/gitlab/instrumentation/redis_base_spec.rb b/spec/lib/gitlab/instrumentation/redis_base_spec.rb
index f9dd0c94c97..656e6ffba05 100644
--- a/spec/lib/gitlab/instrumentation/redis_base_spec.rb
+++ b/spec/lib/gitlab/instrumentation/redis_base_spec.rb
@@ -1,8 +1,10 @@
 # frozen_string_literal: true
 
 require 'spec_helper'
+require 'rspec-parameterized'
 
 RSpec.describe Gitlab::Instrumentation::RedisBase, :request_store do
+  using RSpec::Parameterized::TableSyntax
   let(:instrumentation_class_a) do
     stub_const('InstanceA', Class.new(described_class))
   end
@@ -88,6 +90,44 @@ RSpec.describe Gitlab::Instrumentation::RedisBase, :request_store do
     end
   end
 
+  describe '.increment_cross_slot_request_count' do
+    context 'storage key overlapping' do
+      it 'keys do not overlap across storages' do
+        3.times { instrumentation_class_a.increment_cross_slot_request_count }
+        2.times { instrumentation_class_b.increment_cross_slot_request_count }
+
+        expect(instrumentation_class_a.get_cross_slot_request_count).to eq(3)
+        expect(instrumentation_class_b.get_cross_slot_request_count).to eq(2)
+      end
+
+      it 'increments by the given amount' do
+        instrumentation_class_a.increment_cross_slot_request_count(2)
+        instrumentation_class_a.increment_cross_slot_request_count(3)
+
+        expect(instrumentation_class_a.get_cross_slot_request_count).to eq(5)
+      end
+    end
+  end
+
+  describe '.increment_allowed_cross_slot_request_count' do
+    context 'storage key overlapping' do
+      it 'keys do not overlap across storages' do
+        3.times { instrumentation_class_a.increment_allowed_cross_slot_request_count }
+        2.times { instrumentation_class_b.increment_allowed_cross_slot_request_count }
+
+        expect(instrumentation_class_a.get_allowed_cross_slot_request_count).to eq(3)
+        expect(instrumentation_class_b.get_allowed_cross_slot_request_count).to eq(2)
+      end
+
+      it 'increments by the given amount' do
+        instrumentation_class_a.increment_allowed_cross_slot_request_count(2)
+        instrumentation_class_a.increment_allowed_cross_slot_request_count(3)
+
+        expect(instrumentation_class_a.get_allowed_cross_slot_request_count).to eq(5)
+      end
+    end
+  end
+
   describe '.increment_read_bytes' do
     context 'storage key overlapping' do
       it 'keys do not overlap across storages' do
@@ -130,4 +170,44 @@ RSpec.describe Gitlab::Instrumentation::RedisBase, :request_store do
       end
     end
   end
+
+  describe '.redis_cluster_validate!' do
+    let(:args) { [[:mget, 'foo', 'bar']] }
+
+    before do
+      instrumentation_class_a.enable_redis_cluster_validation
+    end
+
+    context 'Rails environments' do
+      where(:env, :allowed, :should_raise) do
+        'production' | false | false
+        'production' | true | false
+        'staging' | false | false
+        'staging' | true | false
+        'development' | true | false
+        'development' | false | true
+        'test' | true | false
+        'test' | false | true
+      end
+
+      with_them do
+        it do
+          stub_rails_env(env)
+
+          validation = -> { instrumentation_class_a.redis_cluster_validate!(args) }
+          under_test = if allowed
+                         -> { Gitlab::Instrumentation::RedisClusterValidator.allow_cross_slot_commands(&validation) }
+                       else
+                         validation
+                       end
+
+          if should_raise
+            expect(&under_test).to raise_error(::Gitlab::Instrumentation::RedisClusterValidator::CrossSlotError)
+          else
+            expect(&under_test).not_to raise_error
+          end
+        end
+      end
+    end
+  end
 end