From 351371b7f19ac97f4ed85cf2641890d489216da0 Mon Sep 17 00:00:00 2001
From: Dominic Cleal <dominic@cleal.org>
Date: Fri, 27 Jan 2017 11:01:41 +0000
Subject: [PATCH] Deep freeze only strings, arrays and hashes

Freezing a class prevents it from being modified at all, which can
happen in unit tests with mocha or rspec-expections.

Calling a function with `String` as an argument will freeze the entire
class and prevent `allow_any_instance_of(String)...` type expectations
from adding their hooks to the class.

e.g. stdlib:
  https://github.com/puppetlabs/puppetlabs-stdlib/blob/00c881d/spec/functions/is_a_spec.rb#L18
  https://github.com/puppetlabs/puppetlabs-stdlib/blob/00c881d/spec/functions/pw_hash_spec.rb#L57

This whitelists strings, arrays and hashes as safe classes to freeze
when testing functions, others are left thawed to prevent side effects.
Deep freezing is used on arrays/hashes to match Puppet's behaviour of
freezing facts.
---
 .../example/function_example_group.rb         | 22 ++++++++++++++++---
 .../puppet/functions/frozen_array_function.rb |  5 +++++
 .../lib/puppet/functions/frozen_function.rb   |  2 +-
 .../puppet/functions/frozen_hash_function.rb  |  5 +++++
 spec/functions/test_function_spec.rb          |  6 ++++-
 5 files changed, 35 insertions(+), 5 deletions(-)
 create mode 100644 spec/fixtures/modules/test/lib/puppet/functions/frozen_array_function.rb
 create mode 100644 spec/fixtures/modules/test/lib/puppet/functions/frozen_hash_function.rb

diff --git a/lib/rspec-puppet/example/function_example_group.rb b/lib/rspec-puppet/example/function_example_group.rb
index dc3b1bf8..3edbfda7 100644
--- a/lib/rspec-puppet/example/function_example_group.rb
+++ b/lib/rspec-puppet/example/function_example_group.rb
@@ -15,10 +15,8 @@ def initialize(name, func, overrides)
 
       # This method is used by the `run` matcher to trigger the function execution, and provides a uniform interface across all puppet versions.
       def execute(*args)
-        # puppet 4 arguments are immutable
-        args.map(&:freeze)
         Puppet.override(@overrides, "rspec-test scope") do
-          @func.call(@overrides[:global_scope], *args)
+          @func.call(@overrides[:global_scope], *freeze_arg(args))
         end
       end
 
@@ -27,6 +25,24 @@ def call(scope, *args)
         RSpec.deprecate("subject.call", :replacement => "is_expected.to run.with().and_raise_error(), or execute()")
         execute(*args)
       end
+
+      private
+
+      # Facts, keywords, single-quoted strings etc. are usually frozen in Puppet manifests, so freeze arguments to ensure functions are tested
+      # under worst-case conditions.
+      def freeze_arg(arg)
+        case arg
+        when Array
+          arg.each { |a| freeze_arg(a) }
+          arg.freeze
+        when Hash
+          arg.each { |k,v| freeze_arg(k); freeze_arg(v) }
+          arg.freeze
+        when String
+          arg.freeze
+        end
+        arg
+      end
     end
 
     class V3FunctionWrapper
diff --git a/spec/fixtures/modules/test/lib/puppet/functions/frozen_array_function.rb b/spec/fixtures/modules/test/lib/puppet/functions/frozen_array_function.rb
new file mode 100644
index 00000000..6454fc63
--- /dev/null
+++ b/spec/fixtures/modules/test/lib/puppet/functions/frozen_array_function.rb
@@ -0,0 +1,5 @@
+Puppet::Functions.create_function(:frozen_array_function) do
+  def frozen_array_function(value)
+    value.frozen? && value.all?(&:frozen?)
+  end
+end
diff --git a/spec/fixtures/modules/test/lib/puppet/functions/frozen_function.rb b/spec/fixtures/modules/test/lib/puppet/functions/frozen_function.rb
index b67cb15a..41956249 100644
--- a/spec/fixtures/modules/test/lib/puppet/functions/frozen_function.rb
+++ b/spec/fixtures/modules/test/lib/puppet/functions/frozen_function.rb
@@ -1,5 +1,5 @@
 Puppet::Functions.create_function(:frozen_function) do
   def frozen_function(value)
-    value.reverse!
+    value.frozen?
   end
 end
diff --git a/spec/fixtures/modules/test/lib/puppet/functions/frozen_hash_function.rb b/spec/fixtures/modules/test/lib/puppet/functions/frozen_hash_function.rb
new file mode 100644
index 00000000..da641048
--- /dev/null
+++ b/spec/fixtures/modules/test/lib/puppet/functions/frozen_hash_function.rb
@@ -0,0 +1,5 @@
+Puppet::Functions.create_function(:frozen_hash_function) do
+  def frozen_hash_function(value)
+    value.frozen? && value.all? { |k,v| k.frozen? && v.frozen? }
+  end
+end
diff --git a/spec/functions/test_function_spec.rb b/spec/functions/test_function_spec.rb
index 5174b6b2..902c4ad8 100644
--- a/spec/functions/test_function_spec.rb
+++ b/spec/functions/test_function_spec.rb
@@ -6,5 +6,9 @@
 end
 
 describe 'frozen_function', :if => Puppet.version.to_f >= 4.0 do
-  it { is_expected.to run.with_params('foo').and_raise_error(RuntimeError, %r{can't modify frozen}) }
+  it { is_expected.to run.with_params('foo').and_return(true) }
+  it { is_expected.to run.with_params(String).and_return(false) }
+  it { is_expected.to run.with_params(true).and_return(true) }
+  it { is_expected.to run.with_params(['foo']).and_return(true) }
+  it { is_expected.to run.with_params('foo' => 'bar').and_return(true) }
 end