From baf000b07e6faf897b369957a15f2eebe9b19fe7 Mon Sep 17 00:00:00 2001
From: Eric Putnam <putnam.eric@gmail.com>
Date: Thu, 29 Mar 2018 08:48:52 -0700
Subject: [PATCH] (PDOC-229) fix error with return_type and @return

Currently, if the user is documenting a Puppet 4 type custom function and they try to use the return_type method without also including an @return tag, Strings will throw an error. This is caused by the function handler trying to add a type to a return tag on the function object that doesn't exist. To fix the problem, if you are using return_type, the handler will now add an empty return tag first if one does not exist.
---
 .../yard/handlers/ruby/function_handler.rb    |  3 ++
 .../handlers/ruby/function_handler_spec.rb    | 33 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/lib/puppet-strings/yard/handlers/ruby/function_handler.rb b/lib/puppet-strings/yard/handlers/ruby/function_handler.rb
index 52d0da6..ddcd8b2 100644
--- a/lib/puppet-strings/yard/handlers/ruby/function_handler.rb
+++ b/lib/puppet-strings/yard/handlers/ruby/function_handler.rb
@@ -139,6 +139,9 @@ class PuppetStrings::Yard::Handlers::Ruby::FunctionHandler < PuppetStrings::Yard
       next unless DISPATCH_METHOD_NAMES.include?(method_name)
 
       if method_name == 'return_type'
+        # Add a return tag if missing
+        overload_tag.add_tag YARD::Tags::Tag.new(:return, '', 'Any') if overload_tag.tag(:return).nil?
+
         overload_tag.tag(:return).types = [node_as_string(child.parameters[0])]
         next
       end
diff --git a/spec/unit/puppet-strings/yard/handlers/ruby/function_handler_spec.rb b/spec/unit/puppet-strings/yard/handlers/ruby/function_handler_spec.rb
index f141155..d23fe25 100644
--- a/spec/unit/puppet-strings/yard/handlers/ruby/function_handler_spec.rb
+++ b/spec/unit/puppet-strings/yard/handlers/ruby/function_handler_spec.rb
@@ -267,6 +267,39 @@ SOURCE
       end
     end
 
+    describe 'parsing a function using only return_type' do
+      let(:source) { <<-SOURCE
+# An example 4.x function.
+Puppet::Functions.create_function(:foo) do
+  # @param param1 The first parameter.
+  # @param param2 The second parameter.
+  # @param param3 The third parameter.
+  dispatch :foo do
+    param          'Integer',       :param1
+    param          'Any',           :param2
+    optional_param 'Array[String]', :param3
+    return_type 'String'
+  end
+
+  def foo(param1, param2, param3 = nil)
+    "Bar"
+  end
+end
+SOURCE
+      }
+
+      it 'does not throw an error with no @return' do
+        expect { subject }.not_to raise_error NoMethodError
+      end
+
+      it 'contains a return data type' do
+        tags = subject.first.docstring.tags(:return)
+        expect(tags.size).to eq(1)
+        expect(tags[0].name).to be_nil
+        expect(tags[0].types).to eq(['String'])
+      end
+    end
+
     describe 'parsing a function with various dispatch parameters.' do
       let(:source) { <<-SOURCE
 # An example 4.x function.