1 files changed, 81 insertions, 33 deletions
diff --git a/spec/workers/container_registry/migration/guard_worker_spec.rb b/spec/workers/container_registry/migration/guard_worker_spec.rb
index 299d1204af3..c52a3fc5d54 100644
--- a/spec/workers/container_registry/migration/guard_worker_spec.rb
+++ b/spec/workers/container_registry/migration/guard_worker_spec.rb
@@ -25,57 +25,94 @@ RSpec.describe ContainerRegistry::Migration::GuardWorker, :aggregate_failures do
         allow(::Gitlab).to receive(:com?).and_return(true)
       end
 
-      shared_examples 'handling long running migrations' do
+      shared_examples 'handling long running migrations' do |timeout:|
         before do
           allow_next_found_instance_of(ContainerRepository) do |repository|
             allow(repository).to receive(:migration_cancel).and_return(migration_cancel_response)
           end
         end
 
-        context 'migration is canceled' do
-          let(:migration_cancel_response) { { status: :ok } }
-
-          it 'will not abort the migration' do
+        shared_examples 'aborting the migration' do
+          it 'will abort the migration' do
             expect(worker).to receive(:log_extra_metadata_on_done).with(:stale_migrations_count, 1)
             expect(worker).to receive(:log_extra_metadata_on_done).with(:aborted_stale_migrations_count, 1)
             expect(worker).to receive(:log_extra_metadata_on_done).with(:aborted_long_running_migration_ids, [stale_migration.id])
+            expect(ContainerRegistry::Migration).to receive(timeout).and_call_original
 
             expect { subject }
-                .to change(import_skipped_migrations, :count)
+                .to change(import_aborted_migrations, :count).by(1)
+                .and change { stale_migration.reload.migration_state }.to('import_aborted')
+                .and not_change { ongoing_migration.migration_state }
+          end
+
+          context 'registry_migration_guard_thresholds feature flag disabled' do
+            before do
+              stub_feature_flags(registry_migration_guard_thresholds: false)
+            end
+
+            it 'falls back on the hardcoded value' do
+              expect(ContainerRegistry::Migration).not_to receive(:pre_import_timeout)
 
-            expect(stale_migration.reload.migration_state).to eq('import_skipped')
-            expect(stale_migration.reload.migration_skipped_reason).to eq('migration_canceled')
+              expect { subject }
+                .to change { stale_migration.reload.migration_state }.to('import_aborted')
+            end
           end
         end
 
-        context 'migration cancelation fails with an error' do
-          let(:migration_cancel_response) { { status: :error } }
+        context 'migration is canceled' do
+          let(:migration_cancel_response) { { status: :ok } }
 
-          it 'will abort the migration' do
-            expect(worker).to receive(:log_extra_metadata_on_done).with(:stale_migrations_count, 1)
-            expect(worker).to receive(:log_extra_metadata_on_done).with(:aborted_stale_migrations_count, 1)
-            expect(worker).to receive(:log_extra_metadata_on_done).with(:aborted_long_running_migration_ids, [stale_migration.id])
+          before do
+            stub_application_setting(container_registry_import_max_retries: 3)
+          end
 
-            expect { subject }
-                .to change(import_aborted_migrations, :count).by(1)
-                .and change { stale_migration.reload.migration_state }.to('import_aborted')
-                .and not_change { ongoing_migration.migration_state }
+          context 'when the retry limit has been reached' do
+            before do
+              stale_migration.update_column(:migration_retries_count, 2)
+            end
+
+            it 'will not abort the migration' do
+              expect(worker).to receive(:log_extra_metadata_on_done).with(:stale_migrations_count, 1)
+              expect(worker).to receive(:log_extra_metadata_on_done).with(:aborted_stale_migrations_count, 1)
+              expect(worker).to receive(:log_extra_metadata_on_done).with(:aborted_long_running_migration_ids, [stale_migration.id])
+              expect(ContainerRegistry::Migration).to receive(timeout).and_call_original
+
+              expect { subject }
+                  .to change(import_skipped_migrations, :count)
+
+              expect(stale_migration.reload.migration_state).to eq('import_skipped')
+              expect(stale_migration.reload.migration_skipped_reason).to eq('migration_canceled')
+            end
+
+            context 'registry_migration_guard_thresholds feature flag disabled' do
+              before do
+                stub_feature_flags(registry_migration_guard_thresholds: false)
+              end
+
+              it 'falls back on the hardcoded value' do
+                expect(ContainerRegistry::Migration).not_to receive(timeout)
+
+                expect { subject }
+                  .to change { stale_migration.reload.migration_state }.to('import_skipped')
+              end
+            end
+          end
+
+          context 'when the retry limit has not been reached' do
+            it_behaves_like 'aborting the migration'
           end
         end
 
+        context 'migration cancelation fails with an error' do
+          let(:migration_cancel_response) { { status: :error } }
+
+          it_behaves_like 'aborting the migration'
+        end
+
         context 'migration receives bad request with a new status' do
           let(:migration_cancel_response) { { status: :bad_request, migration_state: :import_done } }
 
-          it 'will abort the migration' do
-            expect(worker).to receive(:log_extra_metadata_on_done).with(:stale_migrations_count, 1)
-            expect(worker).to receive(:log_extra_metadata_on_done).with(:aborted_stale_migrations_count, 1)
-            expect(worker).to receive(:log_extra_metadata_on_done).with(:aborted_long_running_migration_ids, [stale_migration.id])
-
-            expect { subject }
-                .to change(import_aborted_migrations, :count).by(1)
-                .and change { stale_migration.reload.migration_state }.to('import_aborted')
-                .and not_change { ongoing_migration.migration_state }
-          end
+          it_behaves_like 'aborting the migration'
         end
       end
 
@@ -96,13 +133,15 @@ RSpec.describe ContainerRegistry::Migration::GuardWorker, :aggregate_failures do
 
       context 'with pre_importing stale migrations' do
         let(:ongoing_migration) { create(:container_repository, :pre_importing) }
-        let(:stale_migration) { create(:container_repository, :pre_importing, migration_pre_import_started_at: 35.minutes.ago) }
+        let(:stale_migration) { create(:container_repository, :pre_importing, migration_pre_import_started_at: 11.minutes.ago) }
         let(:import_status) { 'test' }
 
         before do
           allow_next_instance_of(ContainerRegistry::GitlabApiClient) do |client|
             allow(client).to receive(:import_status).and_return(import_status)
           end
+
+          stub_application_setting(container_registry_pre_import_timeout: 10.minutes.to_i)
         end
 
         it 'will abort the migration' do
@@ -122,13 +161,13 @@ RSpec.describe ContainerRegistry::Migration::GuardWorker, :aggregate_failures do
         context 'the client returns pre_import_in_progress' do
           let(:import_status) { 'pre_import_in_progress' }
 
-          it_behaves_like 'handling long running migrations'
+          it_behaves_like 'handling long running migrations', timeout: :pre_import_timeout
         end
       end
 
       context 'with pre_import_done stale migrations' do
         let(:ongoing_migration) { create(:container_repository, :pre_import_done) }
-        let(:stale_migration) { create(:container_repository, :pre_import_done, migration_pre_import_done_at: 35.minutes.ago) }
+        let(:stale_migration) { create(:container_repository, :pre_import_done, migration_pre_import_done_at: 11.minutes.ago) }
 
         before do
           allow(::ContainerRegistry::Migration).to receive(:max_step_duration).and_return(5.minutes)
@@ -151,13 +190,15 @@ RSpec.describe ContainerRegistry::Migration::GuardWorker, :aggregate_failures do
 
       context 'with importing stale migrations' do
         let(:ongoing_migration) { create(:container_repository, :importing) }
-        let(:stale_migration) { create(:container_repository, :importing, migration_import_started_at: 35.minutes.ago) }
+        let(:stale_migration) { create(:container_repository, :importing, migration_import_started_at: 11.minutes.ago) }
         let(:import_status) { 'test' }
 
         before do
           allow_next_instance_of(ContainerRegistry::GitlabApiClient) do |client|
             allow(client).to receive(:import_status).and_return(import_status)
           end
+
+          stub_application_setting(container_registry_import_timeout: 10.minutes.to_i)
         end
 
         it 'will abort the migration' do
@@ -177,7 +218,7 @@ RSpec.describe ContainerRegistry::Migration::GuardWorker, :aggregate_failures do
         context 'the client returns import_in_progress' do
           let(:import_status) { 'import_in_progress' }
 
-          it_behaves_like 'handling long running migrations'
+          it_behaves_like 'handling long running migrations', timeout: :import_timeout
         end
       end
     end
@@ -195,4 +236,11 @@ RSpec.describe ContainerRegistry::Migration::GuardWorker, :aggregate_failures do
       end
     end
   end
+
+  describe 'worker attributes' do
+    it 'has deduplication set' do
+      expect(described_class.get_deduplicate_strategy).to eq(:until_executed)
+      expect(described_class.get_deduplication_options).to include(ttl: 5.minutes)
+    end
+  end
 end