Fix number of arguments check in flatten()

The function only uses the first argument, so raise an error with
too few arguments *and* with too many arguments.

2 files changed, 5 insertions, 1 deletions

diff --git a/lib/puppet/parser/functions/flatten.rb b/lib/puppet/parser/functions/flatten.rb
index 781da78..a1ed183 100644
--- a/lib/puppet/parser/functions/flatten.rb
+++ b/lib/puppet/parser/functions/flatten.rb
@@ -16,7 +16,7 @@ Would return: ['a','b','c']
   ) do |arguments|
 
     raise(Puppet::ParseError, "flatten(): Wrong number of arguments " +
-      "given (#{arguments.size} for 1)") if arguments.size < 1
+      "given (#{arguments.size} for 1)") if arguments.size != 1
 
     array = arguments[0]
 
diff --git a/spec/unit/puppet/parser/functions/flatten_spec.rb b/spec/unit/puppet/parser/functions/flatten_spec.rb
index d4dfd20..dba7a6b 100755
--- a/spec/unit/puppet/parser/functions/flatten_spec.rb
+++ b/spec/unit/puppet/parser/functions/flatten_spec.rb
@@ -11,6 +11,10 @@ describe "the flatten function" do
     lambda { scope.function_flatten([]) }.should( raise_error(Puppet::ParseError))
   end
 
+  it "should raise a ParseError if there is more than 1 argument" do
+    lambda { scope.function_flatten([[], []]) }.should( raise_error(Puppet::ParseError))
+  end
+
   it "should flatten a complex data structure" do
     result = scope.function_flatten([["a","b",["c",["d","e"],"f","g"]]])
     result.should(eq(["a","b","c","d","e","f","g"]))