1 files changed, 23 insertions, 3 deletions
diff --git a/spec/models/diff_viewer/server_side_spec.rb b/spec/models/diff_viewer/server_side_spec.rb
index 28660b0d4b9..db0814af422 100644
--- a/spec/models/diff_viewer/server_side_spec.rb
+++ b/spec/models/diff_viewer/server_side_spec.rb
@@ -17,10 +17,30 @@ RSpec.describe DiffViewer::ServerSide do
   subject { viewer_class.new(diff_file) }
 
   describe '#prepare!' do
-    it 'loads all diff file data' do
-      expect(Blob).to receive(:lazy).at_least(:twice)
+    before do
+      stub_feature_flags(disable_load_entire_blob_for_diff_viewer: feature_flag_enabled)
+    end
+
+    context 'when the disable_load_entire_blob_for_diff_viewer flag is disabled' do
+      let(:feature_flag_enabled) { false }
 
-      subject.prepare!
+      it 'loads all diff file data' do
+        subject
+        expect(diff_file).to receive_message_chain(:old_blob, :load_all_data!)
+        expect(diff_file).to receive_message_chain(:new_blob, :load_all_data!)
+        subject.prepare!
+      end
+    end
+
+    context 'when the disable_load_entire_blob_for_diff_viewer flag is enabled' do
+      let(:feature_flag_enabled) { true }
+
+      it 'does not load file data' do
+        subject
+        expect(diff_file).not_to receive(:old_blob)
+        expect(diff_file).not_to receive(:new_blob)
+        subject.prepare!
+      end
     end
   end