1 files changed, 58 insertions, 9 deletions

diff --git a/spec/lib/gitlab/utils/strong_memoize_spec.rb b/spec/lib/gitlab/utils/strong_memoize_spec.rb
index 27bfe181ef6..ea8083e7d7f 100644
--- a/spec/lib/gitlab/utils/strong_memoize_spec.rb
+++ b/spec/lib/gitlab/utils/strong_memoize_spec.rb
@@ -3,12 +3,15 @@
 require 'fast_spec_helper'
 require 'rspec-benchmark'
 require 'rspec-parameterized'
+require 'active_support/testing/time_helpers'
 
 RSpec.configure do |config|
   config.include RSpec::Benchmark::Matchers
 end
 
 RSpec.describe Gitlab::Utils::StrongMemoize, feature_category: :shared do
+  include ActiveSupport::Testing::TimeHelpers
+
   let(:klass) do
     strong_memoize_class = described_class
 
@@ -30,6 +33,13 @@ RSpec.describe Gitlab::Utils::StrongMemoize, feature_category: :shared do
         end
       end
 
+      def method_name_with_expiration
+        strong_memoize_with_expiration(:method_name_with_expiration, 1) do
+          trace << value
+          value
+        end
+      end
+
       def method_name_attr
         trace << value
         value
@@ -142,6 +152,43 @@ RSpec.describe Gitlab::Utils::StrongMemoize, feature_category: :shared do
     end
   end
 
+  describe '#strong_memoize_with_expiration' do
+    [nil, false, true, 'value', 0, [0]].each do |value|
+      context "with value #{value}" do
+        let(:value) { value }
+        let(:method_name) { :method_name_with_expiration }
+
+        it_behaves_like 'caching the value'
+
+        it 'raises exception for invalid type as key' do
+          expect { object.strong_memoize_with_expiration(10, 1) { 20 } }.to raise_error /Invalid type of '10'/
+        end
+
+        it 'raises exception for invalid characters in key' do
+          expect { object.strong_memoize_with_expiration(:enabled?, 1) { 20 } }
+            .to raise_error /is not allowed as an instance variable name/
+        end
+      end
+    end
+
+    context 'value memoization test' do
+      let(:value) { 'value' }
+
+      it 'caches the value for specified number of seconds' do
+        object.method_name_with_expiration
+        object.method_name_with_expiration
+
+        expect(object.trace.count).to eq(1)
+
+        travel_to(Time.current + 2.seconds) do
+          object.method_name_with_expiration
+
+          expect(object.trace.count).to eq(2)
+        end
+      end
+    end
+  end
+
   describe '#strong_memoize_with' do
     [nil, false, true, 'value', 0, [0]].each do |value|
       context "with value #{value}" do
@@ -215,19 +262,21 @@ RSpec.describe Gitlab::Utils::StrongMemoize, feature_category: :shared do
 
   describe '.strong_memoize_attr' do
     [nil, false, true, 'value', 0, [0]].each do |value|
-      let(:value) { value }
+      context "with value '#{value}'" do
+        let(:value) { value }
 
-      context "memoized after method definition with value #{value}" do
-        let(:method_name) { :method_name_attr }
+        context 'memoized after method definition' do
+          let(:method_name) { :method_name_attr }
 
-        it_behaves_like 'caching the value'
+          it_behaves_like 'caching the value'
 
-        it 'calls the existing .method_added' do
-          expect(klass.method_added_list).to include(:method_name_attr)
-        end
+          it 'calls the existing .method_added' do
+            expect(klass.method_added_list).to include(:method_name_attr)
+          end
 
-        it 'retains method arity' do
-          expect(klass.instance_method(method_name).arity).to eq(0)
+          it 'retains method arity' do
+            expect(klass.instance_method(method_name).arity).to eq(0)
+          end
         end
       end
     end