1 files changed, 51 insertions, 26 deletions

diff --git a/spec/gitlab_shell_two_factor_recovery_spec.rb b/spec/gitlab_shell_two_factor_recovery_spec.rb
index 872fa85..19999e5 100644
--- a/spec/gitlab_shell_two_factor_recovery_spec.rb
+++ b/spec/gitlab_shell_two_factor_recovery_spec.rb
@@ -10,7 +10,12 @@ describe 'bin/gitlab-shell 2fa_recovery_codes' do
       res.content_type = 'application/json'
       res.status = 200
 
-      key_id = req.query['key_id'] || JSON.parse(req.body)['key_id']
+      key_id = req.query['key_id'] || req.query['user_id']
+
+      unless key_id
+        body = JSON.parse(req.body)
+        key_id = body['key_id'] || body['user_id'].to_s
+      end
 
       if key_id == '100'
         res.body = '{"success":true, "recovery_codes": ["1", "2"]}'
@@ -18,43 +23,63 @@ describe 'bin/gitlab-shell 2fa_recovery_codes' do
         res.body = '{"success":false, "message": "Forbidden!"}'
       end
     end
+
+    server.mount_proc('/api/v4/internal/discover') do |req, res|
+      res.status = 200
+      res.content_type = 'application/json'
+      res.body = '{"id":100, "name": "Some User", "username": "someuser"}'
+    end
   end
 
   shared_examples 'dialog for regenerating recovery keys' do
-    context 'when runs successfully' do
-      let(:cmd) { "#{gitlab_shell_path} key-100" }
+    context 'when the user agrees to regenerate keys' do
+      def verify_successful_regeneration!(cmd)
+        Open3.popen2(env, cmd) do |stdin, stdout|
+          expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
+          expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+
+          stdin.puts('yes')
+
+          expect(stdout.flush.read).to eq(
+            "\nYour two-factor authentication recovery codes are:\n\n" \
+            "1\n2\n\n" \
+            "During sign in, use one of the codes above when prompted for\n" \
+            "your two-factor code. Then, visit your Profile Settings and add\n" \
+            "a new device so you do not lose access to your account again.\n"
+          )
+        end
+      end
+
+      context 'when key is provided' do
+        let(:cmd) { "#{gitlab_shell_path} key-100" }
 
-      context 'when the user agrees to regenerate keys' do
         it 'the recovery keys are regenerated' do
-          Open3.popen2(env, cmd) do |stdin, stdout|
-            expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
-            expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+          verify_successful_regeneration!(cmd)
+        end
+      end
 
-            stdin.puts('yes')
+      context 'when username is provided' do
+        let(:cmd) { "#{gitlab_shell_path} username-someone" }
 
-            expect(stdout.flush.read).to eq(
-              "\nYour two-factor authentication recovery codes are:\n\n" \
-              "1\n2\n\n" \
-              "During sign in, use one of the codes above when prompted for\n" \
-              "your two-factor code. Then, visit your Profile Settings and add\n" \
-              "a new device so you do not lose access to your account again.\n"
-            )
-          end
+        it 'the recovery keys are regenerated' do
+          verify_successful_regeneration!(cmd)
         end
       end
+    end
 
-      context 'when the user disagrees to regenerate keys' do
-        it 'the recovery keys are not regenerated' do
-          Open3.popen2(env, cmd) do |stdin, stdout|
-            expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
-            expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
+    context 'when the user disagrees to regenerate keys' do
+      let(:cmd) { "#{gitlab_shell_path} key-100" }
 
-            stdin.puts('no')
+      it 'the recovery keys are not regenerated' do
+        Open3.popen2(env, cmd) do |stdin, stdout|
+          expect(stdout.gets).to eq("Are you sure you want to generate new two-factor recovery codes?\n")
+          expect(stdout.gets).to eq("Any existing recovery codes you saved will be invalidated. (yes/no)\n")
 
-            expect(stdout.flush.read).to eq(
-              "\nNew recovery codes have *not* been generated. Existing codes will remain valid.\n"
-            )
-          end
+          stdin.puts('no')
+
+          expect(stdout.flush.read).to eq(
+            "\nNew recovery codes have *not* been generated. Existing codes will remain valid.\n"
+          )
         end
       end
     end