1 files changed, 44 insertions, 22 deletions

diff --git a/spec/lib/gitlab/job_waiter_spec.rb b/spec/lib/gitlab/job_waiter_spec.rb
index b000f55e739..e3d7d59df04 100644
--- a/spec/lib/gitlab/job_waiter_spec.rb
+++ b/spec/lib/gitlab/job_waiter_spec.rb
@@ -62,6 +62,10 @@ RSpec.describe Gitlab::JobWaiter, :redis, feature_category: :shared do
 
     before do
       allow_any_instance_of(described_class).to receive(:wait).and_call_original
+      stub_feature_flags(
+        use_primary_and_secondary_stores_for_shared_state: false,
+        use_primary_store_as_default_for_shared_state: false
+      )
     end
 
     it 'returns when all jobs have been completed' do
@@ -83,36 +87,54 @@ RSpec.describe Gitlab::JobWaiter, :redis, feature_category: :shared do
       expect(result).to contain_exactly('a')
     end
 
-    context 'when a label is provided' do
-      let(:waiter) { described_class.new(2, worker_label: 'Foo') }
-      let(:started_total) { double(:started_total) }
-      let(:timeouts_total) { double(:timeouts_total) }
+    context 'when migration is ongoing' do
+      let(:waiter) { described_class.new(3) }
 
-      before do
-        allow(Gitlab::Metrics).to receive(:counter)
-                                    .with(described_class::STARTED_METRIC, anything)
-                                    .and_return(started_total)
+      shared_examples 'returns all jobs' do
+        it 'returns all jobs' do
+          result = nil
+          expect { Timeout.timeout(6) { result = waiter.wait(5) } }.not_to raise_error
 
-        allow(Gitlab::Metrics).to receive(:counter)
-                                    .with(described_class::TIMEOUTS_METRIC, anything)
-                                    .and_return(timeouts_total)
+          expect(result).to contain_exactly('a', 'b', 'c')
+        end
       end
 
-      it 'increments just job_waiter_started_total when all jobs complete' do
-        expect(started_total).to receive(:increment).with(worker: 'Foo')
-        expect(timeouts_total).not_to receive(:increment)
+      context 'when using both stores' do
+        context 'with existing jobs in old store' do
+          before do
+            described_class.notify(waiter.key, 'a')
+            described_class.notify(waiter.key, 'b')
+            described_class.notify(waiter.key, 'c')
+            stub_feature_flags(use_primary_and_secondary_stores_for_shared_state: true)
+          end
 
-        described_class.notify(waiter.key, 'a')
-        described_class.notify(waiter.key, 'b')
+          it_behaves_like 'returns all jobs'
+        end
 
-        expect { Timeout.timeout(1) { waiter.wait(2) } }.not_to raise_error
-      end
+        context 'with jobs in both stores' do
+          before do
+            stub_feature_flags(use_primary_and_secondary_stores_for_shared_state: true)
+            described_class.notify(waiter.key, 'a')
+            described_class.notify(waiter.key, 'b')
+            described_class.notify(waiter.key, 'c')
+          end
 
-      it 'increments job_waiter_started_total and job_waiter_timeouts_total when it times out' do
-        expect(started_total).to receive(:increment).with(worker: 'Foo')
-        expect(timeouts_total).to receive(:increment).with(worker: 'Foo')
+          it_behaves_like 'returns all jobs'
+        end
 
-        expect { Timeout.timeout(2) { waiter.wait(1) } }.not_to raise_error
+        context 'when using primary store as default store' do
+          before do
+            stub_feature_flags(
+              use_primary_and_secondary_stores_for_shared_state: true,
+              use_primary_store_as_default_for_shared_state: true
+            )
+            described_class.notify(waiter.key, 'a')
+            described_class.notify(waiter.key, 'b')
+            described_class.notify(waiter.key, 'c')
+          end
+
+          it_behaves_like 'returns all jobs'
+        end
       end
     end
   end