From 35bf5fd8c93d5052ecf5284ed3194a92cab838d5 Mon Sep 17 00:00:00 2001
From: Martin Foot <martin@modeltwozero.com>
Date: Fri, 21 Feb 2014 14:32:32 +0000
Subject: [PATCH] Allow concat to take non-array second parameters

Also improve and extend concat tests to lock down functionality
---
 README.markdown                               |  5 +++++
 lib/puppet/parser/functions/concat.rb         | 10 +++++++---
 .../puppet/parser/functions/concat_spec.rb    | 19 +++++++++++++++++--
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/README.markdown b/README.markdown
index 1e70f3904..3a6c0cf4c 100644
--- a/README.markdown
+++ b/README.markdown
@@ -145,6 +145,11 @@ Would result in:
 
   ['1','2','3','4','5','6']
 
+  concat(['1','2','3'],'4')
+
+Would result in:
+
+  ['1','2','3','4']
 
 - *Type*: rvalue
 
diff --git a/lib/puppet/parser/functions/concat.rb b/lib/puppet/parser/functions/concat.rb
index c86aa0057..6c8638222 100644
--- a/lib/puppet/parser/functions/concat.rb
+++ b/lib/puppet/parser/functions/concat.rb
@@ -23,12 +23,16 @@ module Puppet::Parser::Functions
     a = arguments[0]
     b = arguments[1]
 
-    # Check that both args are arrays.
-    unless a.is_a?(Array) and b.is_a?(Array)
+    # Check that the first parameter is an array
+    unless a.is_a?(Array)
       raise(Puppet::ParseError, 'concat(): Requires array to work with')
     end
 
-    result = a.concat(b)
+    if b.is_a?(Array)
+      result = a.concat(b)
+    else
+      result = a << b
+    end
 
     return result
   end
diff --git a/spec/unit/puppet/parser/functions/concat_spec.rb b/spec/unit/puppet/parser/functions/concat_spec.rb
index 123188bd4..6e6762096 100644
--- a/spec/unit/puppet/parser/functions/concat_spec.rb
+++ b/spec/unit/puppet/parser/functions/concat_spec.rb
@@ -4,12 +4,27 @@
 describe "the concat function" do
   let(:scope) { PuppetlabsSpec::PuppetInternals.scope }
 
-  it "should raise a ParseError if there is less than 1 arguments" do
-    lambda { scope.function_concat([]) }.should( raise_error(Puppet::ParseError))
+  it "should raise a ParseError if the client does not provide two arguments" do
+    lambda { scope.function_concat([]) }.should(raise_error(Puppet::ParseError))
+  end
+
+  it "should raise a ParseError if the first parameter is not an array" do
+    lambda { scope.function_concat([1, []])}.should(raise_error(Puppet::ParseError))
   end
 
   it "should be able to concat an array" do
     result = scope.function_concat([['1','2','3'],['4','5','6']])
     result.should(eq(['1','2','3','4','5','6']))
   end
+
+  it "should be able to concat a primitive to an array" do
+    result = scope.function_concat([['1','2','3'],'4'])
+    result.should(eq(['1','2','3','4']))
+  end
+
+  it "should not accidentally flatten nested arrays" do
+    result = scope.function_concat([['1','2','3'],[['4','5'],'6']])
+    result.should(eq(['1','2','3',['4','5'],'6']))
+  end
+
 end