1 files changed, 35 insertions, 1 deletions
diff --git a/spec/controllers/projects/merge_requests/diffs_controller_spec.rb b/spec/controllers/projects/merge_requests/diffs_controller_spec.rb
index 4de724fd6d6..23a33d7e0b1 100644
--- a/spec/controllers/projects/merge_requests/diffs_controller_spec.rb
+++ b/spec/controllers/projects/merge_requests/diffs_controller_spec.rb
@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Projects::MergeRequests::DiffsController do
+RSpec.describe Projects::MergeRequests::DiffsController, feature_category: :code_review_workflow do
   include ProjectForksHelper
   include TrackingHelpers
 
@@ -610,5 +610,39 @@ RSpec.describe Projects::MergeRequests::DiffsController do
         go
       end
     end
+
+    context 'when ck param is present' do
+      let(:cache_key) { merge_request.merge_head_diff.id }
+
+      before do
+        create(:merge_request_diff, :merge_head, merge_request: merge_request)
+      end
+
+      it 'sets Cache-Control with max-age' do
+        go(ck: cache_key, diff_head: true)
+
+        expect(response.headers['Cache-Control']).to eq('max-age=86400, private')
+      end
+
+      context 'when diffs_batch_cache_with_max_age feature flag is disabled' do
+        before do
+          stub_feature_flags(diffs_batch_cache_with_max_age: false)
+        end
+
+        it 'does not set Cache-Control with max-age' do
+          go(ck: cache_key, diff_head: true)
+
+          expect(response.headers['Cache-Control']).not_to eq('max-age=86400, private')
+        end
+      end
+
+      context 'when not rendering merge head diff' do
+        it 'does not set Cache-Control with max-age' do
+          go(ck: cache_key, diff_head: false)
+
+          expect(response.headers['Cache-Control']).not_to eq('max-age=86400, private')
+        end
+      end
+    end
   end
 end