3 files changed, 118 insertions, 15 deletions
diff --git a/lib/puppet/network/format_handler.rb b/lib/puppet/network/format_handler.rb
index cf19c02a7..4c9f4e59e 100644
--- a/lib/puppet/network/format_handler.rb
+++ b/lib/puppet/network/format_handler.rb
@@ -63,14 +63,6 @@ module Puppet::Network::FormatHandler
         def supported_formats
             format_handler.formats.collect { |f| format_handler.format(f) }.find_all { |f| f.supported?(self) }.collect { |f| f.name }
         end
-
-        def from_marshal(text)
-            Marshal.load(text)
-        end
-
-        def from_yaml(text)
-            YAML.load(text)
-        end
     end
 
     module InstanceMethods
@@ -83,10 +75,6 @@ module Puppet::Network::FormatHandler
         def support_format?(name)
             self.class.support_format?(name)
         end
-
-        def to_marshal(instance)
-            Marshal.dump(instance)
-        end
     end
 end
 
diff --git a/lib/puppet/network/formats.rb b/lib/puppet/network/formats.rb
index 640a24485..79e2860af 100644
--- a/lib/puppet/network/formats.rb
+++ b/lib/puppet/network/formats.rb
@@ -1,4 +1,54 @@
 require 'puppet/network/format_handler'
 
-Puppet::Network::FormatHandler.create(:yaml, :mime => "text/yaml")
-Puppet::Network::FormatHandler.create(:marshal, :mime => "text/marshal")
+Puppet::Network::FormatHandler.create(:yaml, :mime => "text/yaml") do
+    # Yaml doesn't need the class name; it's serialized.
+    def intern(klass, text)
+        YAML.load(text)
+    end
+
+    # Yaml doesn't need the class name; it's serialized.
+    def intern_multiple(klass, text)
+        YAML.load(text)
+    end
+
+    def render(instance)
+        instance.to_yaml
+    end
+
+    # Yaml monkey-patches Array, so this works.
+    def render_multiple(instances)
+        instances.to_yaml
+    end
+
+    # Everything's supported
+    def supported?(klass)
+        true
+    end
+end
+
+
+Puppet::Network::FormatHandler.create(:marshal, :mime => "text/marshal") do
+    # Yaml doesn't need the class name; it's serialized.
+    def intern(klass, text)
+        Marshal.load(text)
+    end
+
+    # Yaml doesn't need the class name; it's serialized.
+    def intern_multiple(klass, text)
+        Marshal.load(text)
+    end
+
+    def render(instance)
+        Marshal.dump(instance)
+    end
+
+    # Yaml monkey-patches Array, so this works.
+    def render_multiple(instances)
+        Marshal.dump(instances)
+    end
+
+    # Everything's supported
+    def supported?(klass)
+        true
+    end
+end
diff --git a/spec/unit/network/formats.rb b/spec/unit/network/formats.rb
index 94104289d..a7d74e95c 100755
--- a/spec/unit/network/formats.rb
+++ b/spec/unit/network/formats.rb
@@ -10,8 +10,40 @@ describe "Puppet Network Format" do
     end
 
     describe "yaml" do
+        before do
+            @yaml = Puppet::Network::FormatHandler.format(:yaml)
+        end
+
         it "should have its mime type set to text/yaml" do
-            Puppet::Network::FormatHandler.format(:yaml).mime.should == "text/yaml"
+            @yaml.mime.should == "text/yaml"
+        end
+
+        it "should be supported on Strings" do
+            @yaml.should be_supported(String)
+        end
+
+        it "should render by calling 'to_yaml' on the instance" do
+            instance = mock 'instance'
+            instance.expects(:to_yaml).returns "foo"
+            @yaml.render(instance).should == "foo"
+        end
+
+        it "should render multiple instances by calling 'to_yaml' on the array" do
+            instances = [mock('instance')]
+            instances.expects(:to_yaml).returns "foo"
+            @yaml.render_multiple(instances).should == "foo"
+        end
+
+        it "should intern by calling 'YAML.load'" do
+            text = "foo"
+            YAML.expects(:load).with("foo").returns "bar"
+            @yaml.intern(String, text).should == "bar"
+        end
+
+        it "should intern multiples by calling 'YAML.load'" do
+            text = "foo"
+            YAML.expects(:load).with("foo").returns "bar"
+            @yaml.intern_multiple(String, text).should == "bar"
         end
     end
 
@@ -20,8 +52,41 @@ describe "Puppet Network Format" do
     end
 
     describe "marshal" do
+        before do
+            @marshal = Puppet::Network::FormatHandler.format(:marshal)
+        end
+
         it "should have its mime type set to text/marshal" do
             Puppet::Network::FormatHandler.format(:marshal).mime.should == "text/marshal"
         end
+
+        it "should be supported on Strings" do
+            @marshal.should be_supported(String)
+        end
+
+        it "should render by calling 'Marshal.dump' on the instance" do
+            instance = mock 'instance'
+            Marshal.expects(:dump).with(instance).returns "foo"
+            @marshal.render(instance).should == "foo"
+        end
+
+        it "should render multiple instances by calling 'to_marshal' on the array" do
+            instances = [mock('instance')]
+
+            Marshal.expects(:dump).with(instances).returns "foo"
+            @marshal.render_multiple(instances).should == "foo"
+        end
+
+        it "should intern by calling 'Marshal.load'" do
+            text = "foo"
+            Marshal.expects(:load).with("foo").returns "bar"
+            @marshal.intern(String, text).should == "bar"
+        end
+
+        it "should intern multiples by calling 'Marshal.load'" do
+            text = "foo"
+            Marshal.expects(:load).with("foo").returns "bar"
+            @marshal.intern_multiple(String, text).should == "bar"
+        end
     end
 end