(#4655) Allow stage to be set using a default class parameter

For example:
  stage{ pre: before => Stage[main] }
  class someclass ($stage=pre ) { ... }
  class { someclass: }

This transplants adding the edge from the resource to the stage from
the compiler into when the resource is evaluated.  This moves adding
the stage edges to after when the defaults are copied into the
resources, making them available.

Paired-with: Jesse Wolfe <jesse@puppetlabs.com>

6 files changed, 104 insertions, 66 deletions

diff --git a/lib/puppet/parser/compiler.rb b/lib/puppet/parser/compiler.rb
index 6e8e3d26b..98bf3b574 100644
--- a/lib/puppet/parser/compiler.rb
+++ b/lib/puppet/parser/compiler.rb
@@ -56,23 +56,20 @@ class Puppet::Parser::Compiler
     # Note that this will fail if the resource is not unique.
     @catalog.add_resource(resource)
 
+    if resource.type.to_s.downcase != "class" && resource[:stage]
+      raise ArgumentError, "Only classes can set 'stage'; normal resources like #{resource} cannot change run stage"
+    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.
+    # Stages should not be inside of classes.  They are always a
+    # top-level container, regardless of where they appear in the
+    # manifest.
     return if resource.type.to_s.downcase == "stage"
 
+    # This adds a resource to the class it lexically appears in in the
+    # manifest.
     if resource.type.to_s.downcase != "class"
-      raise ArgumentError, "Only classes can set 'stage'; normal resources like #{resource} cannot change run stage" if resource[:stage]
       return @catalog.add_edge(scope.resource, resource)
     end
-
-    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 unless stage.title == :main
-    @catalog.add_edge(stage, resource)
   end
 
   # Do we use nodes found in the code, vs. the external node sources?
diff --git a/lib/puppet/parser/resource.rb b/lib/puppet/parser/resource.rb
index c007d4dbe..b98bc0b80 100644
--- a/lib/puppet/parser/resource.rb
+++ b/lib/puppet/parser/resource.rb
@@ -62,13 +62,28 @@ class Puppet::Parser::Resource < Puppet::Resource
     scope.environment
   end
 
+  # Process the  stage metaparameter for a class.   A containment edge
+  # is drawn from  the class to the stage.   The stage for containment
+  # defaults to main, if none is specified.
+  def add_edge_to_stage
+    unless stage = catalog.resource(:stage, self[:stage] || (scope && scope.resource && scope.resource[:stage]) || :main)
+      raise ArgumentError, "Could not find stage #{self[:stage] || :main} specified by #{self}"
+    end
+
+    self[:stage] ||= stage.title unless stage.title == :main
+    catalog.add_edge(stage, self)
+  end
+
   # Retrieve the associated definition and evaluate it.
   def evaluate
     return if evaluated?
     @evaluated = true
     if klass = resource_type and ! builtin_type?
       finish
-      return klass.evaluate_code(self)
+      evaluated_code = klass.evaluate_code(self)
+      add_edge_to_stage
+
+      return evaluated_code
     elsif builtin?
       devfail "Cannot evaluate a builtin type (#{type})"
     else
diff --git a/lib/puppet/parser/scope.rb b/lib/puppet/parser/scope.rb
index 24f1d01f7..c369f129a 100644
--- a/lib/puppet/parser/scope.rb
+++ b/lib/puppet/parser/scope.rb
@@ -101,7 +101,7 @@ class Puppet::Parser::Scope
 
   # Remove this when rebasing
   def environment
-    compiler.environment
+    compiler ? compiler.environment : nil
   end
 
   # Are we the top scope?
@@ -513,6 +513,6 @@ class Puppet::Parser::Scope
 
   def extend_with_functions_module
     extend Puppet::Parser::Functions.environment_module(Puppet::Node::Environment.root)
-    extend Puppet::Parser::Functions.environment_module(compiler ? environment : nil)
+    extend Puppet::Parser::Functions.environment_module(environment)
   end
 end
diff --git a/spec/unit/parser/compiler_spec.rb b/spec/unit/parser/compiler_spec.rb
index 261cfdec1..4cab3cb65 100755
--- a/spec/unit/parser/compiler_spec.rb
+++ b/spec/unit/parser/compiler_spec.rb
@@ -33,6 +33,14 @@ class CompilerTestResource
 
   def evaluate
   end
+
+  def file
+    "/fake/file/goes/here"
+  end
+
+  def line
+    "42"
+  end
 end
 
 describe Puppet::Parser::Compiler do
@@ -418,52 +426,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 dae22fcaa..6207df36f 100755
--- a/spec/unit/parser/resource_spec.rb
+++ b/spec/unit/parser/resource_spec.rb
@@ -132,9 +132,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
@@ -142,17 +152,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 9895f446b..639512874 100755
--- a/spec/unit/parser/scope_spec.rb
+++ b/spec/unit/parser/scope_spec.rb
@@ -123,7 +123,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
@@ -418,13 +422,15 @@ describe Puppet::Parser::Scope do
       before do
         @scopes = {}
         klass = @scope.known_resource_types.add(Puppet::Resource::Type.new(:hostclass, ""))
-        Puppet::Parser::Resource.new("class", :main, :scope => @scope, :source => mock('source')).evaluate
+        @catalog = Puppet::Resource::Catalog.new
+        @catalog.add_resource(Puppet::Parser::Resource.new("stage", :main, :scope => @scope))
+        Puppet::Parser::Resource.new("class", :main, :scope => @scope, :source => mock('source'), :catalog => @catalog).evaluate
         @scopes[""] = @scope.class_scope(klass)
         @scopes[""].setvar("test", "value")
 
         %w{one one::two one::two::three}.each do |name|
           klass = @scope.known_resource_types.add(Puppet::Resource::Type.new(:hostclass, name))
-          Puppet::Parser::Resource.new("class", name, :scope => @scope, :source => mock('source')).evaluate
+          Puppet::Parser::Resource.new("class", name, :scope => @scope, :source => mock('source'), :catalog => @catalog).evaluate
           @scopes[name] = @scope.class_scope(klass)
           @scopes[name].setvar("test", "value-#{name.sub(/.+::/,'')}")
         end