Minor edits to two_factor_recovery_codes API error catching

2 files changed, 11 insertions, 7 deletions

diff --git a/lib/api/internal.rb b/lib/api/internal.rb
index 5b54c11ef62..6e6efece7c4 100644
--- a/lib/api/internal.rb
+++ b/lib/api/internal.rb
@@ -105,15 +105,19 @@ module API
       post '/two_factor_recovery_codes' do
         status 200
 
-        key = Key.find(params[:key_id])
-        user = key.user
+        key = Key.find_by(id: params[:key_id])
+
+        unless key
+          return { 'success' => false, 'message' => 'Could not find the given key' }
+        end
 
-        # Make sure this isn't a deploy key
-        unless key.type.nil?
+        if key.is_a?(DeployKey)
           return { success: false, message: 'Deploy keys cannot be used to retrieve recovery codes' }
         end
 
-        unless user.present?
+        user = key.user
+
+        unless user
           return { success: false, message: 'Could not find a user for the given key' }
         end
 
diff --git a/spec/requests/api/internal_spec.rb b/spec/requests/api/internal_spec.rb
index 5d06abcfeb3..46d1b868782 100644
--- a/spec/requests/api/internal_spec.rb
+++ b/spec/requests/api/internal_spec.rb
@@ -44,8 +44,8 @@ describe API::API, api: true  do
            secret_token: secret_token,
            key_id: 12345
 
-      expect(response).to have_http_status(404)
-      expect(json_response['message']).to eq('404 Not found')
+      expect(json_response['success']).to be_falsey
+      expect(json_response['message']).to eq('Could not find the given key')
     end
 
     it 'returns an error message when the key is a deploy key' do