Making the Functions module more resilient

We were previously storing the module name with
the environment instances as the key, which meant
if the environment instances were removed, we could
never get those modules again.  Given that the
functions weren't reloaded, this meant the functions
were gone if we ever reloaded the environment.

This makes the Functions environment module resilient across
environment reloads, and it also makes the method work correctly
when handed either an environment or a string.

Signed-off-by: Luke Kanies <luke@puppetlabs.com>
Reviewed-by: Nick Lewis <nick@puppetlabs.com>

2 files changed, 5 insertions, 2 deletions

diff --git a/lib/puppet/parser/functions.rb b/lib/puppet/parser/functions.rb
index e19ac127f..22eee70d7 100644
--- a/lib/puppet/parser/functions.rb
+++ b/lib/puppet/parser/functions.rb
@@ -29,8 +29,11 @@ module Puppet::Parser::Functions
   Environment = Puppet::Node::Environment
 
   def self.environment_module(env = nil)
+    if env and ! env.is_a?(Puppet::Node::Environment)
+      env = Puppet::Node::Environment.new(env)
+    end
     @modules.synchronize {
-      @modules[ env || Environment.current || Environment.root ] ||= Module.new
+      @modules[ (env || Environment.current || Environment.root).name ] ||= Module.new
     }
   end
 
diff --git a/spec/unit/parser/scope_spec.rb b/spec/unit/parser/scope_spec.rb
index 00212a2d5..d3ea8dfa8 100755
--- a/spec/unit/parser/scope_spec.rb
+++ b/spec/unit/parser/scope_spec.rb
@@ -42,7 +42,7 @@ describe Puppet::Parser::Scope do
   end
 
   it "should get its environment from its compiler" do
-    env = stub 'environment'
+    env = Puppet::Node::Environment.new
     compiler = stub 'compiler', :environment => env
     scope = Puppet::Parser::Scope.new :compiler => compiler
     scope.environment.should equal(env)