2 files changed, 25 insertions, 21 deletions

diff --git a/spec/workers/ssh_keys/expired_notification_worker_spec.rb b/spec/workers/ssh_keys/expired_notification_worker_spec.rb
index 249ee404870..109d24f03ab 100644
--- a/spec/workers/ssh_keys/expired_notification_worker_spec.rb
+++ b/spec/workers/ssh_keys/expired_notification_worker_spec.rb
@@ -15,6 +15,20 @@ RSpec.describe SshKeys::ExpiredNotificationWorker, type: :worker do
   describe '#perform' do
     let_it_be(:user) { create(:user) }
 
+    context 'with a large batch' do
+      before do
+        stub_const("SshKeys::ExpiredNotificationWorker::BATCH_SIZE", 5)
+      end
+
+      let_it_be_with_reload(:keys) { create_list(:key, 20, expires_at: 3.days.ago, user: user) }
+
+      it 'updates all keys regardless of batch size' do
+        worker.perform
+
+        expect(keys.pluck(:expiry_notification_delivered_at)).not_to include(nil)
+      end
+    end
+
     context 'with expiring key today' do
       let_it_be_with_reload(:expired_today) { create(:key, expires_at: Time.current, user: user) }
 
@@ -35,24 +49,24 @@ RSpec.describe SshKeys::ExpiredNotificationWorker, type: :worker do
           perform_multiple(worker: worker)
         end
       end
+    end
+
+    context 'when key has expired in the past' do
+      let_it_be(:expired_past) { create(:key, expires_at: 1.day.ago, user: user) }
+
+      it 'does update notified column' do
+        expect { worker.perform }.to change { expired_past.reload.expiry_notification_delivered_at }
+      end
 
-      context 'when feature is not enabled' do
+      context 'when key has already been notified of expiration' do
         before do
-          stub_feature_flags(ssh_key_expiration_email_notification: false)
+          expired_past.update!(expiry_notification_delivered_at: 1.day.ago)
         end
 
         it 'does not update notified column' do
-          expect { worker.perform }.not_to change { expired_today.reload.expiry_notification_delivered_at }
+          expect { worker.perform }.not_to change { expired_past.reload.expiry_notification_delivered_at }
         end
       end
     end
-
-    context 'when key has expired in the past' do
-      let_it_be(:expired_past) { create(:key, expires_at: 1.day.ago, user: user) }
-
-      it 'does not update notified column' do
-        expect { worker.perform }.not_to change { expired_past.reload.expiry_notification_delivered_at }
-      end
-    end
   end
 end
diff --git a/spec/workers/ssh_keys/expiring_soon_notification_worker_spec.rb b/spec/workers/ssh_keys/expiring_soon_notification_worker_spec.rb
index f9276c86cdf..0a1d4a14ad0 100644
--- a/spec/workers/ssh_keys/expiring_soon_notification_worker_spec.rb
+++ b/spec/workers/ssh_keys/expiring_soon_notification_worker_spec.rb
@@ -35,16 +35,6 @@ RSpec.describe SshKeys::ExpiringSoonNotificationWorker, type: :worker do
           perform_multiple(worker: worker)
         end
       end
-
-      context 'when feature is not enabled' do
-        before do
-          stub_feature_flags(ssh_key_expiration_email_notification: false)
-        end
-
-        it 'does not update notified column' do
-          expect { worker.perform }.not_to change { expiring_soon.reload.before_expiry_notification_delivered_at }
-        end
-      end
     end
 
     context 'when key has expired in the past' do