This commit unifies the code paths on which classes are added, alters the 
default
stage to respect the stage of the parent if any, and assures that the resource 
is
notified if its stage is assigned (turning an implicit stage into an explicit 
one).

Signed-off-by: Markus Roberts <[email protected]>
---
 lib/puppet/parser/compiler.rb     |   19 ++++++-------------
 spec/unit/parser/compiler_spec.rb |   13 ++++++++++++-
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/lib/puppet/parser/compiler.rb b/lib/puppet/parser/compiler.rb
index 98cad95..7504b27 100644
--- a/lib/puppet/parser/compiler.rb
+++ b/lib/puppet/parser/compiler.rb
@@ -56,13 +56,10 @@ class Puppet::Parser::Compiler
     # Note that this will fail if the resource is not unique.
     @catalog.add_resource(resource)
 
-    set_container_resource(scope, resource)
-  end
 
-  # Add our container edge.  If we're a class, then we get treated specially - 
we can
-  # control the stage that the class is applied in.  Otherwise, we just
-  # get added to our parent container.
-  def set_container_resource(scope, resource)
+    # Add our container edge.  If we're a class, then we get treated specially 
- we can
+    # control the stage that the class is applied in.  Otherwise, we just
+    # get added to our parent container.
     return if resource.type.to_s.downcase == "stage"
 
     if resource.type.to_s.downcase != "class"
@@ -70,15 +67,14 @@ class Puppet::Parser::Compiler
       return @catalog.add_edge(scope.resource, resource)
     end
 
-    unless stage = @catalog.resource(:stage, resource[:stage] || :main)
+    unless stage = @catalog.resource(:stage, resource[:stage] || (scope && 
scope.resource && scope.resource[:stage]) || :main)
       raise ArgumentError, "Could not find stage #{resource[:stage] || :main} 
specified by #{resource}"
     end
 
+    resource[:stage] ||= stage.title
     @catalog.add_edge(stage, resource)
   end
 
-  private :set_container_resource
-
   # Do we use nodes found in the code, vs. the external node sources?
   def ast_nodes?
     known_resource_types.nodes?
@@ -289,10 +285,7 @@ class Puppet::Parser::Compiler
     @main_resource = Puppet::Parser::Resource.new("class", :main, :scope => 
@topscope, :source => @main)
     @topscope.resource = @main_resource
 
-    @resources << @main_resource
-    @catalog.add_resource(@main_resource)
-
-    set_container_resource(@topscope, @main_resource)
+    add_resource(@topscope, @main_resource)
 
     @main_resource.evaluate
   end
diff --git a/spec/unit/parser/compiler_spec.rb 
b/spec/unit/parser/compiler_spec.rb
index e8c06dd..22d52f2 100755
--- a/spec/unit/parser/compiler_spec.rb
+++ b/spec/unit/parser/compiler_spec.rb
@@ -430,7 +430,18 @@ describe Puppet::Parser::Compiler do
       lambda { @compiler.add_resource(@scope, resource) }.should 
raise_error(ArgumentError)
     end
 
-    it "should add edges from the class resources to the main stage if no 
stage is specified" do
+    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)
-- 
1.7.0.4

-- 
You received this message because you are subscribed to the Google Groups 
"Puppet Developers" group.
To post to this group, send email to [email protected].
To unsubscribe from this group, send email to 
[email protected].
For more options, visit this group at 
http://groups.google.com/group/puppet-dev?hl=en.

Reply via email to