fixes issue by adding extensions to definitions

- if the given status code was not definied as response

Author: LeFnord

lib/grape-swagger/doc_methods/extensions.rb

@@ -4,6 +4,7 @@ class Extensions
       class << self
         def add(path, definitions, route)
           @route = route
+
           description = route.settings[:description]
           add_extension_to(path[method], extension(description)) if description && extended?(description, :x)
 
@@ -20,29 +21,32 @@ def add_extensions_to_definition(settings, path, definitions)
           def_extension = extension(settings, :x_def)
 
           if def_extension[:x_def].is_a?(Array)
-            def_extension[:x_def].each do |extension|
-              next unless extension.key?(:for)
-              status = extension.delete(:for)
-              definition = find_definition(status, path)
-              add_extension_to(definitions[definition], x_def: extension)
-            end
+            def_extension[:x_def].each { |extension| setup_definition(extension, path, definitions) }
           else
-            return unless def_extension[:x_def].key?(:for)
-            status = def_extension[:x_def].delete(:for)
-            definition = find_definition(status, path)
-            add_extension_to(definitions[definition], def_extension)
+            setup_definition(def_extension[:x_def], path, definitions)
           end
         end
 
+        def setup_definition(def_extension, path, definitions)
+          return unless def_extension.key?(:for)
+          status = def_extension[:for]
+
+          definition = find_definition(status, path)
+          add_extension_to(definitions[definition], x_def: def_extension)
+        end
+
         def find_definition(status, path)
           response = path[method][:responses][status]
+          return if response.nil?
 
-          response[:schema]['$ref'].split('/').last
+          return response[:schema]['$ref'].split('/').last if response[:schema].key?('$ref')
+          return response[:schema]['items']['$ref'].split('/').last if response[:schema].key?('items')
         end
 
         def add_extension_to(part, extensions)
+          return if part.nil?
           concatenate(extensions).each do |key, value|
-            part[key] = value
+            part[key] = value unless key.start_with?('x-for')
           end
         end
 

spec/swagger_v2/api_swagger_v2_extensions_spec.rb

@@ -1,6 +1,6 @@
 require 'spec_helper'
 
-describe 'extension' do
+describe 'extensions' do
   include_context "#{MODEL_PARSER} swagger example"
 
   before :all do
@@ -40,6 +40,14 @@ class ExtensionsApi < Grape::API
 
         route_setting :x_def, [{ for: 422, other: 'stuff' }, { for: 200, some: 'stuff' }]
 
+        desc 'This returns something with extension on definition level',
+             success: Entities::OtherItem
+        get '/non_existend_status_definitions_extension' do
+          { 'declared_params' => declared(params) }
+        end
+
+        route_setting :x_def, [{ for: 422, other: 'stuff' }, { for: 200, some: 'stuff' }]
+
         desc 'This returns something with extension on definition level',
              success: Entities::OtherItem,
              failure: [{ code: 422, message: 'NotFound', model: Entities::SecondApiError }]
@@ -106,4 +114,17 @@ def app
       expect(subject['definitions']['SecondApiError']['x-other']).to eql 'stuff'
     end
   end
+
+  describe 'extension on definition level' do
+    subject do
+      get '/swagger_doc/non_existend_status_definitions_extension'
+      JSON.parse(last_response.body)
+    end
+
+    specify do
+      expect(subject['definitions'].length).to eql 1
+      expect(subject['definitions']['OtherItem']).to include 'x-some'
+      expect(subject['definitions']['OtherItem']['x-some']).to eql 'stuff'
+    end
+  end
 end