Adding JSON support to Catalogs

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

4 files changed, 267 insertions, 0 deletions

diff --git a/lib/puppet/relationship.rb b/lib/puppet/relationship.rb
index 4c44adba7..8efebf1e3 100644
--- a/lib/puppet/relationship.rb
+++ b/lib/puppet/relationship.rb
@@ -6,10 +6,13 @@
 # subscriptions are permanent associations determining how different
 # objects react to an event
 
+require 'puppet/util/json'
+
 # This is Puppet's class for modeling edges in its configuration graph.
 # It used to be a subclass of GRATR::Edge, but that class has weird hash
 # overrides that dramatically slow down the graphing.
 class Puppet::Relationship
+    extend Puppet::Util::Json
     attr_accessor :source, :target, :callback
 
     attr_reader :event
diff --git a/lib/puppet/resource/catalog.rb b/lib/puppet/resource/catalog.rb
index eb42ff690..68e6d7de5 100644
--- a/lib/puppet/resource/catalog.rb
+++ b/lib/puppet/resource/catalog.rb
@@ -4,6 +4,7 @@ require 'puppet/simple_graph'
 require 'puppet/transaction'
 
 require 'puppet/util/cacher'
+require 'puppet/util/json'
 
 require 'puppet/util/tagging'
 
@@ -18,6 +19,7 @@ class Puppet::Resource::Catalog < Puppet::SimpleGraph
     indirects :catalog, :terminus_class => :compiler
 
     include Puppet::Util::Tagging
+    extend Puppet::Util::Json
     include Puppet::Util::Cacher::Expirer
 
     # The host name this is a catalog for.
@@ -388,6 +390,69 @@ class Puppet::Resource::Catalog < Puppet::SimpleGraph
         @resource_table.keys
     end
 
+    def self.from_json(data)
+        result = new(data['name'])
+
+        if tags = data['tags'] 
+            result.tag(*tags)
+        end
+
+        if version = data['version'] 
+            result.version = version
+        end
+
+        if resources = data['resources']
+            resources.each do |res|
+                resource_from_json(result, res)
+            end
+        end
+
+        if edges = data['edges']
+            edges.each do |edge|
+                edge_from_json(result, edge)
+            end
+        end
+
+        result
+    end
+
+    def self.edge_from_json(result, edge)
+        # If no json_class information was presented, we manually find
+        # the class.
+        edge = Puppet::Relationship.from_json(edge) if edge.is_a?(Hash)
+        unless source = result.resource(edge.source)
+            raise ArgumentError, "Could not convert from json: Could not find relationship source '%s'" % source
+        end
+        edge.source = source
+
+        unless target = result.resource(edge.target)
+            raise ArgumentError, "Could not convert from json: Could not find relationship target '%s'" % target
+        end
+        edge.target = target
+
+        result.add_edge(edge)
+    end
+
+    def self.resource_from_json(result, res)
+        # If no json_class information was presented, we manually find
+        # the class.
+        res = Puppet::Resource.from_json(res) if res.is_a?(Hash)
+        result.add_resource(res)
+    end
+
+    def to_json(*args)
+        {
+            'json_class' => 'Puppet::Resource::Catalog',
+            'data' => {
+                'tags' => tags,
+                'name' => name,
+                'version' => version,
+                'resources' => vertices.to_json(*args),
+                'edges' => edges.to_json(*args)
+            }
+        }.to_json(*args)
+    end
+
     # Convert our catalog into a RAL catalog.
     def to_ral
         to_catalog :to_ral
diff --git a/spec/unit/relationship.rb b/spec/unit/relationship.rb
index 6aa11ad9f..fd7e0aaf7 100755
--- a/spec/unit/relationship.rb
+++ b/spec/unit/relationship.rb
@@ -72,6 +72,16 @@ describe Puppet::Relationship, " when initializing" do
         @edge.callback.should == :foo
         @edge.event.should == :bar
     end
+
+    it "should accept events specified as strings" do
+        @edge = Puppet::Relationship.new(:a, :b, "event" => :NONE)
+        @edge.event.should == :NONE
+    end
+
+    it "should accept callbacks specified as strings" do
+        @edge = Puppet::Relationship.new(:a, :b, "callback" => :foo)
+        @edge.callback.should == :foo
+    end
 end
 
 describe Puppet::Relationship, " when matching edges with no specified event" do
@@ -227,6 +237,10 @@ describe Puppet::Relationship, "when converting from json" do
         Puppet::Relationship.expects(:new).with { |*args| yield args }
     end
 
+    it "should be extended with the JSON utility module" do
+        Puppet::Relationship.metaclass.ancestors.should be_include(Puppet::Util::Json)
+    end
+
     # LAK:NOTE For all of these tests, we convert back to the edge so we can
     # trap the actual data structure then.
     it "should pass the source in as the first argument" do
diff --git a/spec/unit/resource/catalog.rb b/spec/unit/resource/catalog.rb
index cf6a87461..2f4476a2b 100755
--- a/spec/unit/resource/catalog.rb
+++ b/spec/unit/resource/catalog.rb
@@ -826,3 +826,188 @@ describe Puppet::Resource::Catalog, "when compiling" do
         end
     end
 end
+
+describe Puppet::Resource::Catalog, "when converting to json" do
+    confine "Missing 'json' library" => Puppet.features.json?
+
+    before do
+        @catalog = Puppet::Resource::Catalog.new("myhost")
+    end
+
+    def json_output_should
+        @catalog.class.expects(:json_create).with { |hash| yield hash }
+    end
+
+    # LAK:NOTE For all of these tests, we convert back to the resource so we can
+    # trap the actual data structure then.
+    it "should set its json_class to 'Puppet::Resource::Catalog'" do
+        json_output_should { |hash| hash['json_class'] == "Puppet::Resource::Catalog" }
+
+        JSON.parse @catalog.to_json
+    end
+
+    it "should set its data as a hash" do
+        json_output_should { |hash| hash['data'].is_a?(Hash) }
+        JSON.parse @catalog.to_json
+    end
+
+    [:name, :version, :tags].each do |param|
+        it "should set its #{param} to the #{param} of the resource" do
+            @catalog.send(param.to_s + "=", "testing") unless @catalog.send(param)
+
+            json_output_should { |hash| hash['data'][param.to_s] == @catalog.send(param) }
+            JSON.parse @catalog.to_json
+        end
+    end
+
+    it "should convert its resources to a JSON-encoded array and store it as the 'resources' data" do
+        one = stub 'one', :to_json => '"one_resource"', :ref => "Foo[one]"
+        two = stub 'two', :to_json => '"two_resource"', :ref => "Foo[two]"
+
+        @catalog.add_resource(one)
+        @catalog.add_resource(two)
+
+        # TODO this should really guarantee sort order
+        json_output_should { |hash| JSON.parse(hash['data']['resources']).sort == ["one_resource", "two_resource"].sort }
+        JSON.parse @catalog.to_json
+    end
+
+    it "should convert its edges to a JSON-encoded array and store it as the 'edges' data" do
+        one = stub 'one', :to_json => '"one_resource"', :ref => 'Foo[one]'
+        two = stub 'two', :to_json => '"two_resource"', :ref => 'Foo[two]'
+        three = stub 'three', :to_json => '"three_resource"', :ref => 'Foo[three]'
+
+        @catalog.add_edge(one, two)
+        @catalog.add_edge(two, three)
+
+        @catalog.edge(one, two).expects(:to_json).returns '"one_two_json"'
+        @catalog.edge(two, three).expects(:to_json).returns '"two_three_json"'
+
+        json_output_should { |hash| JSON.parse(hash['data']['edges']).sort == %w{one_two_json two_three_json}.sort }
+        JSON.parse @catalog.to_json
+    end
+end
+
+describe Puppet::Resource::Catalog, "when converting from json" do
+    confine "Missing 'json' library" => Puppet.features.json?
+
+    def json_result_should
+        Puppet::Resource::Catalog.expects(:new).with { |hash| yield hash }
+    end
+
+    before do
+        @data = {
+            'name' => "myhost"
+        }
+        @json = {
+            'json_class' => 'Puppet::Resource::Catalog',
+            'data' => @data
+        }
+
+        @catalog = Puppet::Resource::Catalog.new("myhost")
+        Puppet::Resource::Catalog.stubs(:new).returns @catalog
+    end
+
+    it "should be extended with the JSON utility module" do
+        Puppet::Resource::Catalog.metaclass.ancestors.should be_include(Puppet::Util::Json)
+    end
+
+    it "should create it with the provided name" do
+        Puppet::Resource::Catalog.expects(:new).with('myhost').returns @catalog
+        JSON.parse @json.to_json
+    end
+
+    it "should set the provided version on the catalog if one is set" do
+        @data['version'] = 50
+        @catalog.expects(:version=).with(@data['version'])
+
+        JSON.parse @json.to_json
+    end
+
+    it "should set any provided tags on the catalog" do
+        @data['tags'] = %w{one two}
+        @catalog.expects(:tag).with("one", "two")
+
+        JSON.parse @json.to_json
+    end
+
+    it 'should convert the resources list into resources and add each of them' do
+        @data['resources'] = [Puppet::Resource.new(:file, "/foo"), Puppet::Resource.new(:file, "/bar")]
+
+        @catalog.expects(:add_resource).times(2).with { |res| res.type == "File" }
+
+        JSON.parse @json.to_json
+    end
+
+    it 'should convert resources even if they do not include "json_class" information' do
+        @data['resources'] = [Puppet::Resource.new(:file, "/foo")]
+
+        @data['resources'][0].expects(:to_json).returns "{\"title\":\"\\/foo\",\"tags\":[\"file\"],\"type\":\"File\"}"
+
+        @catalog.expects(:add_resource).with { |res| res.type == "File" }
+
+        JSON.parse @json.to_json
+    end
+
+    it 'should convert the edges list into edges and add each of them' do
+        one = Puppet::Relationship.new("osource", "otarget", :event => "one", :callback => "refresh")
+        two = Puppet::Relationship.new("tsource", "ttarget", :event => "two", :callback => "refresh")
+
+        @data['edges'] = [one, two]
+
+        @catalog.stubs(:resource).returns("eh")
+
+        @catalog.expects(:add_edge).with { |edge| edge.event == "one" }
+        @catalog.expects(:add_edge).with { |edge| edge.event == "two" }
+
+        JSON.parse @json.to_json
+    end
+
+    it "should be able to convert relationships that do not include 'json_class' information" do
+        one = Puppet::Relationship.new("osource", "otarget", :event => "one", :callback => "refresh")
+        one.expects(:to_json).returns "{\"event\":\"one\",\"callback\":\"refresh\",\"source\":\"osource\",\"target\":\"otarget\"}"
+
+        @data['edges'] = [one]
+
+        @catalog.stubs(:resource).returns("eh")
+
+        @catalog.expects(:add_edge).with { |edge| edge.event == "one" }
+
+        JSON.parse @json.to_json
+    end
+
+    it "should set the source and target for each edge to the actual resource" do
+        edge = Puppet::Relationship.new("source", "target")
+
+        @data['edges'] = [edge]
+
+        @catalog.expects(:resource).with("source").returns("source_resource")
+        @catalog.expects(:resource).with("target").returns("target_resource")
+
+        @catalog.expects(:add_edge).with { |edge| edge.source == "source_resource" and edge.target == "target_resource" }
+
+        JSON.parse @json.to_json
+    end
+
+    it "should fail if the source resource cannot be found" do
+        edge = Puppet::Relationship.new("source", "target")
+
+        @data['edges'] = [edge]
+
+        @catalog.expects(:resource).with("source").returns(nil)
+        @catalog.stubs(:resource).with("target").returns("target_resource")
+
+        lambda { JSON.parse @json.to_json }.should raise_error(ArgumentError)
+    end
+
+    it "should fail if the target resource cannot be found" do
+        edge = Puppet::Relationship.new("source", "target")
+
+        @data['edges'] = [edge]
+
+        @catalog.stubs(:resource).with("source").returns("source_resource")
+        @catalog.expects(:resource).with("target").returns(nil)
+
+        lambda { JSON.parse @json.to_json }.should raise_error(ArgumentError)
+    end
+end