2 files changed, 48 insertions, 11 deletions

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 == ""