1 files changed, 63 insertions, 9 deletions

diff --git a/spec/tasks/gitlab/shell_rake_spec.rb b/spec/tasks/gitlab/shell_rake_spec.rb
index 195859eac70..30f512205f9 100644
--- a/spec/tasks/gitlab/shell_rake_spec.rb
+++ b/spec/tasks/gitlab/shell_rake_spec.rb
@@ -24,21 +24,75 @@ RSpec.describe 'gitlab:shell rake tasks', :silence_stdout do
   end
 
   describe 'setup task' do
-    it 'writes authorized keys into the file' do
-      allow(Gitlab::CurrentSettings).to receive(:authorized_keys_enabled?).and_return(true)
-      stub_env('force', 'yes')
+    let!(:auth_key) { create(:key) }
+    let!(:auth_and_signing_key) { create(:key, usage_type: :auth_and_signing) }
 
-      auth_key = create(:key)
-      auth_and_signing_key = create(:key, usage_type: :auth_and_signing)
+    before do
       create(:key, usage_type: :signing)
 
-      expect_next_instance_of(Gitlab::AuthorizedKeys) do |instance|
-        expect(instance).to receive(:batch_add_keys).once do |keys|
-          expect(keys).to match_array([auth_key, auth_and_signing_key])
+      allow(Gitlab::CurrentSettings).to receive(:authorized_keys_enabled?).and_return(write_to_authorized_keys)
+    end
+
+    context 'when "Write to authorized keys" is enabled' do
+      let(:write_to_authorized_keys) { true }
+
+      before do
+        stub_env('force', force)
+      end
+
+      context 'when "force" is not set' do
+        let(:force) { nil }
+
+        context 'when the user answers "yes"' do
+          it 'writes authorized keys into the file' do
+            allow(main_object).to receive(:ask_to_continue)
+
+            expect_next_instance_of(Gitlab::AuthorizedKeys) do |instance|
+              expect(instance).to receive(:batch_add_keys).once do |keys|
+                expect(keys).to match_array([auth_key, auth_and_signing_key])
+              end
+            end
+
+            run_rake_task('gitlab:shell:setup')
+          end
+        end
+
+        context 'when the user answers "no"' do
+          it 'does not write authorized keys into the file' do
+            allow(main_object).to receive(:ask_to_continue).and_raise(Gitlab::TaskAbortedByUserError)
+
+            expect(Gitlab::AuthorizedKeys).not_to receive(:new)
+
+            expect do
+              run_rake_task('gitlab:shell:setup')
+            end.to raise_error(SystemExit)
+          end
+        end
+      end
+
+      context 'when "force" is set to "yes"' do
+        let(:force) { 'yes' }
+
+        it 'writes authorized keys into the file' do
+          expect_next_instance_of(Gitlab::AuthorizedKeys) do |instance|
+            expect(instance).to receive(:batch_add_keys).once do |keys|
+              expect(keys).to match_array([auth_key, auth_and_signing_key])
+            end
+          end
+
+          run_rake_task('gitlab:shell:setup')
         end
       end
+    end
+
+    context 'when "Write to authorized keys" is disabled' do
+      let(:write_to_authorized_keys) { false }
 
-      run_rake_task('gitlab:shell:setup')
+      it 'does not write authorized keys into the file' do
+        expect(Gitlab::AuthorizedKeys).not_to receive(:new)
+
+        run_rake_task('gitlab:shell:setup')
+      end
     end
   end
 end