4 files changed, 70 insertions, 14 deletions

diff --git a/lib/puppet/parser/functions.rb b/lib/puppet/parser/functions.rb
index 38bb2eb42..4179cb81d 100644
--- a/lib/puppet/parser/functions.rb
+++ b/lib/puppet/parser/functions.rb
@@ -2,10 +2,12 @@ require 'puppet/util/autoload'
 require 'puppet/parser/scope'
 
 # A module for managing parser functions.  Each specified function
-# becomes an instance method on the Scope class.
+# is added to a central module that then gets included into the Scope
+# class.
 module Puppet::Parser::Functions
 
     @functions = {}
+    @modules = {}
 
     class << self
         include Puppet::Util
@@ -22,6 +24,10 @@ module Puppet::Parser::Functions
         @autoloader
     end
 
+    def self.environment_module(env = nil)
+        @module ||= Module.new
+    end
+
     # Create a new function type.
     def self.newfunction(name, options = {}, &block)
         name = symbolize(name)
@@ -43,7 +49,7 @@ module Puppet::Parser::Functions
         end
 
         fname = "function_" + name.to_s
-        Puppet::Parser::Scope.send(:define_method, fname, &block)
+        environment_module.send(:define_method, fname, &block)
 
         # Someday we'll support specifying an arity, but for now, nope
         #@functions[name] = {:arity => arity, :type => ftype}
@@ -64,7 +70,7 @@ module Puppet::Parser::Functions
         @functions.delete(name)
 
         fname = "function_" + name.to_s
-        Puppet::Parser::Scope.send(:remove_method, fname)
+        environment_module.send(:remove_method, fname)
     end
 
     # Determine if a given name is a function
diff --git a/lib/puppet/parser/scope.rb b/lib/puppet/parser/scope.rb
index 19d79e6d0..991e12302 100644
--- a/lib/puppet/parser/scope.rb
+++ b/lib/puppet/parser/scope.rb
@@ -105,6 +105,11 @@ class Puppet::Parser::Scope
         end
     end
 
+    # Remove this when rebasing
+    def environment
+        compiler.environment
+    end
+
     # Are we the top scope?
     def topscope?
         @level == 1
@@ -141,6 +146,8 @@ class Puppet::Parser::Scope
             end
         }
 
+        extend_with_functions_module()
+
         @tags = []
 
         # The symbol table for this scope.  This is where we store variables.
@@ -494,4 +501,10 @@ class Puppet::Parser::Scope
             setvar("#{i+1}", m, :file => file, :line => line, :ephemeral => true)
         end
     end
+
+    private
+
+    def extend_with_functions_module
+        extend Puppet::Parser::Functions.environment_module(compiler ? environment : nil)
+    end
 end
diff --git a/spec/unit/parser/functions.rb b/spec/unit/parser/functions.rb
index fe449139d..f605052b5 100644
--- a/spec/unit/parser/functions.rb
+++ b/spec/unit/parser/functions.rb
@@ -4,9 +4,6 @@ require File.dirname(__FILE__) + '/../../spec_helper'
 
 describe Puppet::Parser::Functions do
 
-    before(:each) do
-    end
-
     after(:each) do
         # Rationale:
         # our various tests will almost all register to Pupet::Parser::Functions
@@ -23,33 +20,51 @@ describe Puppet::Parser::Functions do
         end
     end
 
+    it "should have a method for returning an environment-specific module" do
+        Puppet::Parser::Functions.environment_module("myenv").should be_instance_of(Module)
+    end
+
+    it "should use the current default environment if no environment is provided" do
+        Puppet::Parser::Functions.environment_module().should be_instance_of(Module)
+    end
+
     describe "when calling newfunction" do
-        it "should create the function in the scope class" do
-            Puppet::Parser::Scope.expects(:define_method).with { |name,block| name == "function_name" }
+        before do
+            @module = Module.new
+            Puppet::Parser::Functions.stubs(:environment_module).returns @module
+        end
+
+        it "should create the function in the environment module" do
+            @module.expects(:define_method).with { |name,block| name == "function_name" }
 
             Puppet::Parser::Functions.newfunction("name", :type => :rvalue)
         end
 
         it "should raise an error if the function already exists" do
-            Puppet::Parser::Scope.expects(:define_method).with { |name,block| name == "function_name" }.once
+            @module.expects(:define_method).with { |name,block| name == "function_name" }.once
             Puppet::Parser::Functions.newfunction("name", :type => :rvalue)
 
             lambda { Puppet::Parser::Functions.newfunction("name", :type => :rvalue) }.should raise_error
         end
 
         it "should raise an error if the function type is not correct" do
-            Puppet::Parser::Scope.expects(:define_method).with { |name,block| name == "function_name" }.never
+            @module.expects(:define_method).with { |name,block| name == "function_name" }.never
 
             lambda { Puppet::Parser::Functions.newfunction("name", :type => :unknown) }.should raise_error
         end
     end
 
     describe "when calling rmfunction" do
+        before do
+            @module = Module.new
+            Puppet::Parser::Functions.stubs(:environment_module).returns @module
+        end
+
         it "should remove the function in the scope class" do
-            Puppet::Parser::Scope.expects(:define_method).with { |name,block| name == "function_name" }
+            @module.expects(:define_method).with { |name,block| name == "function_name" }
             Puppet::Parser::Functions.newfunction("name", :type => :rvalue)
 
-            Puppet::Parser::Scope.expects(:remove_method).with("function_name").once
+            @module.expects(:remove_method).with("function_name").once
 
             Puppet::Parser::Functions.rmfunction("name")
         end
@@ -60,6 +75,10 @@ describe Puppet::Parser::Functions do
     end
 
     describe "when calling function to test function existance" do
+        before do
+            @module = Module.new
+            Puppet::Parser::Functions.stubs(:environment_module).returns @module
+        end
 
         it "should return false if the function doesn't exist" do
             Puppet::Parser::Functions.autoloader.stubs(:load)
@@ -67,8 +86,8 @@ describe Puppet::Parser::Functions do
             Puppet::Parser::Functions.function("name").should be_false
         end
 
-        it "should return it's name if the function exists" do
-            Puppet::Parser::Scope.expects(:define_method).with { |name,block| name == "function_name" }
+        it "should return its name if the function exists" do
+            @module.expects(:define_method).with { |name,block| name == "function_name" }
             Puppet::Parser::Functions.newfunction("name", :type => :rvalue)
 
             Puppet::Parser::Functions.function("name").should == "function_name"
diff --git a/spec/unit/parser/scope.rb b/spec/unit/parser/scope.rb
index c030f2552..7093279b6 100755
--- a/spec/unit/parser/scope.rb
+++ b/spec/unit/parser/scope.rb
@@ -58,6 +58,24 @@ describe Puppet::Parser::Scope do
         Puppet::Parser::Scope.ancestors.should include(Puppet::Resource::TypeCollectionHelper)
     end
 
+    describe "when initializing" do
+        it "should extend itself with its environment's Functions module" do
+            env = Puppet::Node::Environment.new("myenv")
+            compiler = stub 'compiler', :environment => env
+            mod = Module.new
+            Puppet::Parser::Functions.expects(:environment_module).with(env).returns mod
+
+            Puppet::Parser::Scope.new(:compiler => compiler).metaclass.ancestors.should be_include(mod)
+        end
+
+        it "should extend itself with the default Functions module if it has no environment" do
+            mod = Module.new
+            Puppet::Parser::Functions.expects(:environment_module).with(nil).returns mod
+
+            Puppet::Parser::Scope.new().metaclass.ancestors.should be_include(mod)
+        end
+    end
+
     describe "when looking up a variable" do
         it "should default to an empty string" do
             @scope.lookupvar("var").should == ""