3 files changed, 74 insertions, 50 deletions
diff --git a/spec/unit/parser/compiler_spec.rb b/spec/unit/parser/compiler_spec.rb
index ced760b76..48aeb98dd 100755
--- a/spec/unit/parser/compiler_spec.rb
+++ b/spec/unit/parser/compiler_spec.rb
@@ -32,6 +32,14 @@ class CompilerTestResource
 
   def evaluate
   end
+
+  def file
+    "/fake/file/goes/here"
+  end
+
+  def line
+    "42"
+  end
 end
 
 describe Puppet::Parser::Compiler do
@@ -413,52 +421,6 @@ describe Puppet::Parser::Compiler do
       @compiler.catalog.should be_edge(@scope.resource, resource)
     end
 
-    it "should add an edge to any specified stage for class resources" do
-      other_stage = resource(:stage, "other")
-      @compiler.add_resource(@scope, other_stage)
-      resource = resource(:class, "foo")
-      resource[:stage] = 'other'
-
-      @compiler.add_resource(@scope, resource)
-
-      @compiler.catalog.edge?(other_stage, resource).should be_true
-    end
-
-    it "should fail if a non-class resource attempts to set a stage" do
-      other_stage = resource(:stage, "other")
-      @compiler.add_resource(@scope, other_stage)
-      resource = resource(:file, "foo")
-      resource[:stage] = 'other'
-
-      lambda { @compiler.add_resource(@scope, resource) }.should raise_error(ArgumentError)
-    end
-
-    it "should fail if an unknown stage is specified" do
-      resource = resource(:class, "foo")
-      resource[:stage] = 'other'
-
-      lambda { @compiler.add_resource(@scope, resource) }.should raise_error(ArgumentError)
-    end
-
-    it "should add edges from the class resources to the parent's stage if no stage is specified" do
-      main      = @compiler.catalog.resource(:stage, :main)
-      foo_stage = resource(:stage, :foo_stage)
-      @compiler.add_resource(@scope, foo_stage)
-      resource = resource(:class, "foo")
-      @scope.stubs(:resource).returns(:stage => :foo_stage)
-      @compiler.add_resource(@scope, resource)
-
-      @compiler.catalog.should be_edge(foo_stage, resource)
-    end
-
-    it "should add edges from top-level class resources to the main stage if no stage is specified" do
-      main = @compiler.catalog.resource(:stage, :main)
-      resource = resource(:class, "foo")
-      @compiler.add_resource(@scope, resource)
-
-      @compiler.catalog.should be_edge(main, resource)
-    end
-
     it "should not add non-class resources that don't specify a stage to the 'main' stage" do
       main = @compiler.catalog.resource(:stage, :main)
       resource = resource(:file, "foo")
diff --git a/spec/unit/parser/resource_spec.rb b/spec/unit/parser/resource_spec.rb
index b03c18e5f..0d9cba60b 100755
--- a/spec/unit/parser/resource_spec.rb
+++ b/spec/unit/parser/resource_spec.rb
@@ -131,9 +131,19 @@ describe Puppet::Parser::Resource do
   end
 
   describe "when evaluating" do
+    before do
+      @node = Puppet::Node.new "test-node"
+      @compiler = Puppet::Parser::Compiler.new @node
+      @catalog = Puppet::Resource::Catalog.new
+      source = stub('source')
+      source.stubs(:module_name)
+      @scope = Puppet::Parser::Scope.new(:compiler => @compiler, :source => source)
+      @catalog.add_resource(Puppet::Parser::Resource.new("stage", :main, :scope => @scope))
+    end
+
     it "should evaluate the associated AST definition" do
       definition = newdefine "mydefine"
-      res = Puppet::Parser::Resource.new("mydefine", "whatever", :scope => @scope, :source => @source)
+      res = Puppet::Parser::Resource.new("mydefine", "whatever", :scope => @scope, :source => @source, :catalog => @catalog)
       definition.expects(:evaluate_code).with(res)
 
       res.evaluate
@@ -141,17 +151,65 @@ describe Puppet::Parser::Resource do
 
     it "should evaluate the associated AST class" do
       @class = newclass "myclass"
-      res = Puppet::Parser::Resource.new("class", "myclass", :scope => @scope, :source => @source)
+      res = Puppet::Parser::Resource.new("class", "myclass", :scope => @scope, :source => @source, :catalog => @catalog)
       @class.expects(:evaluate_code).with(res)
       res.evaluate
     end
 
     it "should evaluate the associated AST node" do
       nodedef = newnode("mynode")
-      res = Puppet::Parser::Resource.new("node", "mynode", :scope => @scope, :source => @source)
+      res = Puppet::Parser::Resource.new("node", "mynode", :scope => @scope, :source => @source, :catalog => @catalog)
       nodedef.expects(:evaluate_code).with(res)
       res.evaluate
     end
+
+    it "should add an edge to any specified stage for class resources" do
+      @compiler.known_resource_types.add Puppet::Resource::Type.new(:hostclass, "foo", '')
+
+      other_stage = Puppet::Parser::Resource.new(:stage, "other", :scope => @scope, :catalog => @catalog)
+      @compiler.add_resource(@scope, other_stage)
+      resource = Puppet::Parser::Resource.new(:class, "foo", :scope => @scope, :catalog => @catalog)
+      resource[:stage] = 'other'
+      @compiler.add_resource(@scope, resource)
+
+      resource.evaluate
+
+      @compiler.catalog.edge?(other_stage, resource).should be_true
+    end
+
+    it "should fail if an unknown stage is specified" do
+      @compiler.known_resource_types.add Puppet::Resource::Type.new(:hostclass, "foo", '')
+
+      resource = Puppet::Parser::Resource.new(:class, "foo", :scope => @scope, :catalog => @catalog)
+      resource[:stage] = 'other'
+
+      lambda { resource.evaluate }.should raise_error(ArgumentError, /Could not find stage other specified by/)
+    end
+
+    it "should add edges from the class resources to the parent's stage if no stage is specified" do
+      main      = @compiler.catalog.resource(:stage, :main)
+      foo_stage = Puppet::Parser::Resource.new(:stage, :foo_stage, :scope => @scope, :catalog => @catalog)
+      @compiler.add_resource(@scope, foo_stage)
+      @compiler.known_resource_types.add Puppet::Resource::Type.new(:hostclass, "foo", '')
+      resource = Puppet::Parser::Resource.new(:class, "foo", :scope => @scope, :catalog => @catalog)
+      resource[:stage] = 'foo_stage'
+      @compiler.add_resource(@scope, resource)
+
+      resource.evaluate
+
+      @compiler.catalog.should be_edge(foo_stage, resource)
+    end
+
+    it "should add edges from top-level class resources to the main stage if no stage is specified" do
+      main = @compiler.catalog.resource(:stage, :main)
+      @compiler.known_resource_types.add Puppet::Resource::Type.new(:hostclass, "foo", '')
+      resource = Puppet::Parser::Resource.new(:class, "foo", :scope => @scope, :catalog => @catalog)
+      @compiler.add_resource(@scope, resource)
+
+      resource.evaluate
+
+      @compiler.catalog.should be_edge(main, resource)
+    end
   end
 
   describe "when finishing" do
diff --git a/spec/unit/parser/scope_spec.rb b/spec/unit/parser/scope_spec.rb
index bf4d1e29e..5308856ed 100755
--- a/spec/unit/parser/scope_spec.rb
+++ b/spec/unit/parser/scope_spec.rb
@@ -121,7 +121,11 @@ describe Puppet::Parser::Scope do
 
       def create_class_scope(name)
         klass = newclass(name)
-        Puppet::Parser::Resource.new("class", name, :scope => @scope, :source => mock('source')).evaluate
+
+        catalog = Puppet::Resource::Catalog.new
+        catalog.add_resource(Puppet::Parser::Resource.new("stage", :main, :scope => Puppet::Parser::Scope.new))
+
+        Puppet::Parser::Resource.new("class", name, :scope => @scope, :source => mock('source'), :catalog => catalog).evaluate
 
         @scope.class_scope(klass)
       end