Moving classfile-writing to the Catalog

This work was done by the Agent class before,
but it's really related to the catalog, so that's
where it is now.

Signed-off-by: Luke Kanies <luke@madstop.com>

4 files changed, 45 insertions, 25 deletions

diff --git a/lib/puppet/agent.rb b/lib/puppet/agent.rb
index f0ee101bf..5cb6f0019 100644
--- a/lib/puppet/agent.rb
+++ b/lib/puppet/agent.rb
@@ -138,6 +138,8 @@ class Puppet::Agent
         return nil unless result
 
         result.retrieval_duration = duration
+        result.host_config = true
+        result.write_class_file
         return result
     end
 
@@ -188,25 +190,6 @@ class Puppet::Agent
         lockfile.locked?
     end
 
-    # Store the classes in the classfile, but only if we're not local.
-    def setclasses(ary)
-        if @local
-            return
-        end
-        unless ary and ary.length > 0
-            Puppet.info "No classes to store"
-            return
-        end
-        begin
-            ::File.open(Puppet[:classfile], "w") { |f|
-                f.puts ary.join("\n")
-            }
-        rescue => detail
-            Puppet.err "Could not create class file %s: %s" %
-                [Puppet[:classfile], detail]
-        end
-    end
-
     private
 
     def self.timeout
diff --git a/lib/puppet/resource/catalog.rb b/lib/puppet/resource/catalog.rb
index 0692b5ddd..8f013e7db 100644
--- a/lib/puppet/resource/catalog.rb
+++ b/lib/puppet/resource/catalog.rb
@@ -393,6 +393,17 @@ class Puppet::Resource::Catalog < Puppet::SimpleGraph
         to_catalog :to_resource
     end
 
+    # Store the classes in the classfile.
+    def write_class_file
+        begin
+            ::File.open(Puppet[:classfile], "w") do |f|
+                f.puts classes.join("\n")
+            end
+        rescue => detail
+            Puppet.err "Could not create class file %s: %s" % [Puppet[:classfile], detail]
+        end
+    end
+
     # Produce the graph files if requested.
     def write_graph(name)
         # We only want to graph the main host catalog.
diff --git a/spec/unit/agent.rb b/spec/unit/agent.rb
index c204fa8fb..922577cb6 100755
--- a/spec/unit/agent.rb
+++ b/spec/unit/agent.rb
@@ -142,7 +142,7 @@ describe Puppet::Agent, "when retrieving a catalog" do
         Puppet.settings.stubs(:use).returns(true)
         @agent = Puppet::Agent.new
 
-        @catalog = stub 'catalog', :retrieval_duration= => nil
+        @catalog = Puppet::Resource::Catalog.new
     end
 
     it "should use the Catalog class to get its catalog" do
@@ -189,17 +189,28 @@ describe Puppet::Agent, "when retrieving a catalog" do
     it "should record the retrieval time with the catalog" do
         @agent.expects(:thinmark).yields.then.returns 10
 
-        catalog = mock 'catalog'
-        Puppet::Resource::Catalog.expects(:get).returns catalog
+        Puppet::Resource::Catalog.expects(:get).returns @catalog
 
-        catalog.expects(:retrieval_duration=).with 10
+        @catalog.expects(:retrieval_duration=).with 10
 
         @agent.retrieve_catalog
     end
 
-    it "should update the class file with the classes contained within the catalog"
+    it "should write the catalog's class file" do
+        @catalog.expects(:write_class_file)
+
+        Puppet::Resource::Catalog.expects(:get).returns @catalog
+
+        @agent.retrieve_catalog
+    end
 
-    it "should mark the catalog as a host catalog"
+    it "should mark the catalog as a host catalog" do
+        @catalog.expects(:host_config=).with true
+
+        Puppet::Resource::Catalog.expects(:get).returns @catalog
+
+        @agent.retrieve_catalog
+    end
 
     it "should return nil if there is an error while retrieving the catalog" do
         Puppet::Resource::Catalog.expects(:get).raises "eh"
diff --git a/spec/unit/resource/catalog.rb b/spec/unit/resource/catalog.rb
index c46915f86..bc526a70d 100755
--- a/spec/unit/resource/catalog.rb
+++ b/spec/unit/resource/catalog.rb
@@ -20,6 +20,21 @@ describe Puppet::Resource::Catalog, "when compiling" do
         @catalog.should_not be_expired(Time.now)
     end
 
+    it "should be able to write its list of classes to the class file" do
+        @catalog = Puppet::Resource::Catalog.new("host")
+
+        @catalog.add_class "foo", "bar"
+
+        Puppet.settings.expects(:value).with(:classfile).returns "/class/file"
+
+        fh = mock 'filehandle'
+        File.expects(:open).with("/class/file", "w").yields fh
+
+        fh.expects(:puts).with "foo\nbar"
+
+        @catalog.write_class_file
+    end
+
     describe "when compiling" do
         it "should accept tags" do
             config = Puppet::Resource::Catalog.new("mynode")