1 files changed, 53 insertions, 25 deletions
diff --git a/spec/lib/gitlab/cache/request_store_wrap_spec.rb b/spec/lib/gitlab/cache/request_store_wrap_spec.rb
index 6a95239066b..d63d958900a 100644
--- a/spec/lib/gitlab/cache/request_store_wrap_spec.rb
+++ b/spec/lib/gitlab/cache/request_store_wrap_spec.rb
@@ -5,20 +5,17 @@ describe Gitlab::Cache::RequestStoreWrap, :request_store do
     Class.new do
       extend Gitlab::Cache::RequestStoreWrap
 
-      attr_accessor :id, :name, :result
+      attr_accessor :id, :name, :result, :extra
 
       def self.name
         'ExpensiveAlgorithm'
       end
 
-      def initialize(id, name, result)
+      def initialize(id, name, result, extra = nil)
         self.id = id
         self.name = name
         self.result = result
-      end
-
-      request_store_wrap_key do
-        [id, name]
+        self.extra = nil
       end
 
       request_store_wrap def compute(arg)
@@ -28,6 +25,11 @@ describe Gitlab::Cache::RequestStoreWrap, :request_store do
       request_store_wrap def repute(arg)
         result << arg
       end
+
+      def dispute(arg)
+        result << arg
+      end
+      request_store_wrap(:dispute) { extra }
     end
   end
 
@@ -50,24 +52,6 @@ describe Gitlab::Cache::RequestStoreWrap, :request_store do
       expect(algorithm.result).to eq(result)
     end
 
-    it 'computes twice for the different keys, id' do
-      algorithm.compute(true)
-      algorithm.id = 'ad'
-      result = algorithm.compute(true)
-
-      expect(result).to eq([true, true])
-      expect(algorithm.result).to eq(result)
-    end
-
-    it 'computes twice for the different keys, name' do
-      algorithm.compute(true)
-      algorithm.name = 'same'
-      result = algorithm.compute(true)
-
-      expect(result).to eq([true, true])
-      expect(algorithm.result).to eq(result)
-    end
-
     it 'computes twice for the different class name' do
       algorithm.compute(true)
       allow(klass).to receive(:name).and_return('CheapAlgo')
@@ -95,6 +79,42 @@ describe Gitlab::Cache::RequestStoreWrap, :request_store do
       expect(result).to eq([true, true])
       expect(algorithm.result).to eq(result)
     end
+
+    context 'when request_store_wrap_key is provided' do
+      before do
+        klass.request_store_wrap_key do
+          [id, name]
+        end
+      end
+
+      it 'computes twice for the different keys, id' do
+        algorithm.compute(true)
+        algorithm.id = 'ad'
+        result = algorithm.compute(true)
+
+        expect(result).to eq([true, true])
+        expect(algorithm.result).to eq(result)
+      end
+
+      it 'computes twice for the different keys, name' do
+        algorithm.compute(true)
+        algorithm.name = 'same'
+        result = algorithm.compute(true)
+
+        expect(result).to eq([true, true])
+        expect(algorithm.result).to eq(result)
+      end
+
+      it 'uses extra method cache key if provided' do
+        algorithm.dispute(true) # miss
+        algorithm.extra = true
+        algorithm.dispute(true) # miss
+        result = algorithm.dispute(true) # hit
+
+        expect(result).to eq([true, true])
+        expect(algorithm.result).to eq(result)
+      end
+    end
   end
 
   context 'when RequestStore is inactive' do
@@ -102,10 +122,18 @@ describe Gitlab::Cache::RequestStoreWrap, :request_store do
       RequestStore.end!
     end
 
-    it 'computes twice even if everything is the same' do
+    it 'computes only once if it is the same instance for the same key' do
       algorithm.compute(true)
       result = algorithm.compute(true)
 
+      expect(result).to eq([true])
+      expect(algorithm.result).to eq(result)
+    end
+
+    it 'computes twice for different instances even if keys are the same' do
+      algorithm.compute(true)
+      result = klass.new('id', 'name', algorithm.result).compute(true)
+
       expect(result).to eq([true, true])
       expect(algorithm.result).to eq(result)
     end