Add specs for load balancing calls in Ci::Trace

The methods involved are private methods so I added specs under their
most relevant public methods, based on a46c789f

lib/gitlab/ci/trace.rb

@@ -14,6 +14,8 @@ module Gitlab
       UPDATE_FREQUENCY_DEFAULT = 60.seconds
       UPDATE_FREQUENCY_WHEN_BEING_WATCHED = 3.seconds
 
+      LOAD_BALANCING_STICKING_NAMESPACE = 'ci/build/trace'
+
       ArchiveError = Class.new(StandardError)
       AlreadyArchivedError = Class.new(StandardError)
       LockedError = Class.new(StandardError)
@@ -299,7 +301,7 @@ module Gitlab
       def destroy_stream(build)
         if consistent_archived_trace?(build)
           ::Gitlab::Database::LoadBalancing::Sticking
-            .stick('ci/build/trace', build.id)
+            .stick(LOAD_BALANCING_STICKING_NAMESPACE, build.id)
         end
 
         yield
@@ -308,7 +310,7 @@ module Gitlab
       def read_trace_artifact(build)
         if consistent_archived_trace?(build)
           ::Gitlab::Database::LoadBalancing::Sticking
-            .unstick_or_continue_sticking('ci/build/trace', build.id)
+            .unstick_or_continue_sticking(LOAD_BALANCING_STICKING_NAMESPACE, build.id)
         end
 
         yield

spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb

@@ -29,6 +29,34 @@ RSpec.shared_examples 'common trace features' do
     end
   end
 
+  describe '#read' do
+    context 'gitlab_ci_archived_trace_consistent_reads feature flag enabled' do
+      before do
+        stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project)
+      end
+
+      it 'calls ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking' do
+        expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:unstick_or_continue_sticking)
+          .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id)
+          .and_call_original
+
+        trace.read { |stream| stream }
+      end
+    end
+
+    context 'gitlab_ci_archived_trace_consistent_reads feature flag disabled' do
+      before do
+        stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false)
+      end
+
+      it 'does not call ::Gitlab::Database::LoadBalancing::Sticking.unstick_or_continue_sticking' do
+        expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
+
+        trace.read { |stream| stream }
+      end
+    end
+  end
+
   describe '#extract_coverage' do
     let(:regex) { '\(\d+.\d+\%\) covered' }
 
@@ -790,6 +818,32 @@ RSpec.shared_examples 'trace with enabled live trace feature' do
           .to eq(described_class.hexdigest(build.job_artifacts_trace.file.path))
         expect(build.job_artifacts_trace.file_sha256).to eq(src_checksum)
       end
+
+      context 'gitlab_ci_archived_trace_consistent_reads feature flag enabled' do
+        before do
+          stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: build.project)
+        end
+
+        it 'calls ::Gitlab::Database::LoadBalancing::Sticking.stick' do
+          expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
+            .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, build.id)
+            .and_call_original
+
+          subject
+        end
+      end
+
+      context 'gitlab_ci_archived_trace_consistent_reads feature flag disabled' do
+        before do
+          stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false)
+        end
+
+        it 'does not call ::Gitlab::Database::LoadBalancing::Sticking.stick' do
+          expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:stick)
+
+          subject
+        end
+      end
     end
 
     shared_examples 'source trace in ChunkedIO stays intact' do |error:|