1 files changed, 35 insertions, 0 deletions
diff --git a/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb b/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb
index deb22de9160..9e30564b437 100644
--- a/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb
+++ b/spec/lib/gitlab/import_export/json/streaming_serializer_spec.rb
@@ -156,6 +156,41 @@ RSpec.describe Gitlab::ImportExport::Json::StreamingSerializer do
         subject.execute
       end
     end
+
+    describe 'load balancing' do
+      context 'when feature flag load_balancing_for_export_workers is enabled' do
+        before do
+          stub_feature_flags(load_balancing_for_export_workers: true)
+        end
+
+        context 'when enabled', :db_load_balancing do
+          it 'reads from replica' do
+            expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
+
+            subject.execute
+          end
+        end
+
+        context 'when disabled' do
+          it 'reads from primary' do
+            allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
+            expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_replicas_for_read_queries)
+
+            subject.execute
+          end
+        end
+      end
+
+      context 'when feature flag load_balancing_for_export_workers is disabled' do
+        it 'reads from primary' do
+          stub_feature_flags(load_balancing_for_export_workers: false)
+
+          expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_replicas_for_read_queries)
+
+          subject.execute
+        end
+      end
+    end
   end
 
   describe '.batch_size' do