Extract plugin<->provides mapping to custom class

Code to map a plugin to the attributes it provides or look up a plugin
by provided attributes was located in different parts of the code base
where it operated directly on the underlying data structure. Extracting
to a class lets us access the information we want using domain specific
terminology so things are easier to understand.

5 files changed, 125 insertions, 72 deletions

diff --git a/lib/ohai/dsl/plugin.rb b/lib/ohai/dsl/plugin.rb
index 023f511a..cbfb85e8 100644
--- a/lib/ohai/dsl/plugin.rb
+++ b/lib/ohai/dsl/plugin.rb
@@ -109,12 +109,16 @@ module Ohai
 
       def initialize(controller, source)
         @controller = controller
-        @attributes = controller.attributes
         @data = controller.data
         @source = source
         @has_run = false
       end
 
+      # TODO: rename
+      def attributes
+        @controller.attributes
+      end
+
       def run
         @has_run = true
         run_plugin
@@ -227,19 +231,7 @@ module Ohai
         end
 
         def provides(*paths)
-          paths.each do |path|
-            parts = path.split("/")
-            a = @attributes
-            unless parts.length == 0
-              parts.shift if parts[0].length == 0
-              parts.each do |part|
-                a[part] ||= Mash.new
-                a = a[part]
-              end
-            end
-            a[:_plugins] ||= []
-            a[:_plugins] << self
-          end
+          attributes.set_providers_for(self, paths)
         end
 
         def require_plugin(*args)
diff --git a/lib/ohai/loader.rb b/lib/ohai/loader.rb
index 6428f969..aea67183 100644
--- a/lib/ohai/loader.rb
+++ b/lib/ohai/loader.rb
@@ -68,21 +68,7 @@ module Ohai
 
     def collect_provides(plugin)
       plugin_provides = plugin.class.provides_attrs
-      
-      plugin_provides.each do |attr|
-        parts = attr.split('/')
-        a = @attributes
-        unless parts.length == 0
-          parts.shift if parts[0].length == 0
-          parts.each do |part|
-            a[part] ||= Mash.new
-            a = a[part]
-          end
-        end
-
-        a[:_plugins] ||= []
-        a[:_plugins] << plugin
-      end
+      @attributes.set_providers_for(plugin, plugin_provides)
     end
 
   end
diff --git a/lib/ohai/runner.rb b/lib/ohai/runner.rb
index c16f0f54..82fe4880 100644
--- a/lib/ohai/runner.rb
+++ b/lib/ohai/runner.rb
@@ -32,42 +32,33 @@ module Ohai
     # true, then this plugin and its dependencies will be run even if
     # they have been run before.
     def run_plugin(plugin, force = false)
+      unless plugin.kind_of?(Ohai::DSL::Plugin)
+        raise ArgumentError, "Invalid plugin #{plugin} (must be an Ohai::DSL::Plugin or subclass)"
+      end
       visited = [plugin]
       while !visited.empty?
-        p = visited.pop
+        next_plugin = visited.pop
 
-        next if p.has_run? unless force
+        next if next_plugin.has_run? unless force
 
-        if visited.include?(p)
+        if visited.include?(next_plugin)
           raise Ohai::Exceptions::DependencyCycle, "Dependency cycle detected. Please refer to the following plugins: #{get_cycle(visited, p).join(", ") }"
         end
 
-        dependency_providers = fetch_plugins(p.dependencies)
-        dependency_providers.delete_if { |plugin| (!force && plugin.has_run?) || plugin.eql?(p) }
+        dependency_providers = fetch_plugins(next_plugin.dependencies)
+        dependency_providers.delete_if { |dep_plugin| (!force && dep_plugin.has_run?) || dep_plugin.eql?(next_plugin) }
 
         if dependency_providers.empty?
-          @safe_run ? p.safe_run : p.run
+          @safe_run ? next_plugin.safe_run : next_plugin.run
         else
-          visited << p << dependency_providers.first
+          visited << next_plugin << dependency_providers.first
         end
       end
     end
 
     # returns a list of plugins which provide the given attributes
     def fetch_plugins(attributes)
-      plugins = []
-      attributes.each do |attribute|
-        attrs = @attributes
-        parts = attribute.split('/')
-        parts.each do |part|
-          next if part == Ohai::Mixin::OS.collect_os
-          raise Ohai::Exceptions::AttributeNotFound, "Cannot find plugin providing attribute \'#{attribute}\'" unless attrs[part]
-          attrs = attrs[part]
-        end
-        plugins << attrs[:_plugins]
-        plugins.flatten!
-      end
-      plugins.uniq
+      @attributes.find_providers_for(attributes)
     end
 
     # given a list of plugins and the first plugin in the cycle,
diff --git a/lib/ohai/system.rb b/lib/ohai/system.rb
index 02d41387..b3dfeaf3 100644
--- a/lib/ohai/system.rb
+++ b/lib/ohai/system.rb
@@ -30,6 +30,86 @@ require 'mixlib/shellout'
 require 'yajl'
 
 module Ohai
+
+  class ProvidesMap
+
+    attr_reader :map
+
+    def initialize
+      @map = Mash.new
+    end
+
+    def set_providers_for(plugin, provided_attributes)
+      unless plugin.kind_of?(Ohai::DSL::Plugin)
+        raise ArgumentError, "set_providers_for only accepts Ohai Plugin classes (got: #{plugin})"
+      end
+      provided_attributes.each do |attr|
+        parts = attr.split('/')
+        a = map
+        unless parts.length == 0
+          parts.shift if parts[0].length == 0
+          parts.each do |part|
+            a[part] ||= Mash.new
+            a = a[part]
+          end
+        end
+
+        a[:_plugins] ||= []
+        a[:_plugins] << plugin
+      end
+    end
+
+    def find_providers_for(attributes)
+      plugins = []
+      attributes.each do |attribute|
+        attrs = map
+        parts = attribute.split('/')
+        parts.each do |part|
+          next if part == Ohai::Mixin::OS.collect_os
+          raise Ohai::Exceptions::AttributeNotFound, "Cannot find plugin providing attribute \'#{attribute}\'" unless attrs[part]
+          attrs = attrs[part]
+        end
+        plugins << attrs[:_plugins]
+        plugins.flatten!
+      end
+      plugins.uniq
+    end
+
+
+    def all_plugins
+      collected = []
+      collect_plugins_in(map, collected).uniq
+    end
+
+    # TODO: change this to has_provider? or something domain specific.
+    def has_key?(key)
+      @map.has_key?(key)
+    end
+
+    # TODO: refactor out direct access to the map (mostly only occurs in test code)
+    def [](key)
+      @map[key]
+    end
+
+    # TODO: refactor out direct access to the map (mostly only occurs in test code)
+    def []=(key, value)
+      @map[key] = value
+    end
+
+    private
+
+    def collect_plugins_in(provides_map, collected)
+      provides_map.keys.each do |plugin|
+        if plugin.eql?("_plugins")
+          collected.concat(provides_map[plugin])
+        else
+          collect_plugins_in(provides_map[plugin], collected)
+        end
+      end
+      collected
+    end
+  end
+
   class System
     attr_accessor :data
     attr_reader :attributes
@@ -38,7 +118,8 @@ module Ohai
 
     def initialize
       @data = Mash.new
-      @attributes = Mash.new
+      @attributes = ProvidesMap.new
+
       @hints = Hash.new
       @v6_dependency_solver = Hash.new
       @plugin_path = ""
@@ -92,7 +173,8 @@ module Ohai
       end
 
       # collect and run version 7 plugins
-      plugins = collect_plugins(@attributes)
+      plugins = @attributes.all_plugins
+
       begin
         plugins.each { |plugin| @runner.run_plugin(plugin, force) }
       rescue Ohai::Exceptions::AttributeNotFound, Ohai::Exceptions::DependencyCycle => e
@@ -105,6 +187,7 @@ module Ohai
     def collect_plugins(plugins)
       collected = []
       if plugins.is_a?(Mash)
+        # TODO: remove this branch
         plugins.keys.each do |plugin|
           if plugin.eql?("_plugins")
             collected << plugins[plugin]
diff --git a/spec/unit/system_spec.rb b/spec/unit/system_spec.rb
index 0cdfab52..fc3f323b 100644
--- a/spec/unit/system_spec.rb
+++ b/spec/unit/system_spec.rb
@@ -29,8 +29,8 @@ describe "Ohai::System" do
       @ohai.should be_a_kind_of(Ohai::System)
     end
 
-    it "should set @attributes to a Mash" do
-      @ohai.attributes.should be_a_kind_of(Mash)
+    it "should set @attributes to a ProvidesMap" do
+      @ohai.attributes.should be_a_kind_of(Ohai::ProvidesMap)
     end
 
     it "should set @v6_dependency_solver to a Hash" do
@@ -128,7 +128,7 @@ describe "Ohai::System" do
           @ohai = Ohai::System.new
           klass = Ohai.plugin(:Empty) { }
           plugin = klass.new(@ohai, "/tmp/plugins/empty.rb")
-          @ohai.stub(:collect_plugins).and_return([plugin])
+          @ohai.attributes.should_receive(:all_plugins).and_return([plugin])
         end
 
         describe "when AttributeNotFound is received" do
@@ -165,7 +165,7 @@ describe "Ohai::System" do
           @plugins << klass.new(@ohai, "")
         end
 
-        @ohai.stub(:collect_plugins).and_return(@plugins)
+        @ohai.attributes.should_receive(:all_plugins).and_return(@plugins)
       end
 
       it "should run each plugin once from Ohai::System" do
@@ -253,19 +253,20 @@ EOF
       end
     end
 
+    # TODO: Tests ProvidesMap not System
     it "should find all the plugins providing attributes" do
-      a = @ohai.attributes
-      a[:zero] = Mash.new
-      a[:zero][:_plugins] = [@plugins[0]]
-      a[:one] = Mash.new
-      a[:one][:_plugins] = [@plugins[1]]
-      a[:one][:two] = Mash.new
-      a[:one][:two][:_plugins] = [@plugins[2]]
-      a[:stub] = Mash.new
-      a[:stub][:three] = Mash.new
-      a[:stub][:three][:_plugins] = [@plugins[3]]
-
-      providers = @ohai.collect_plugins(@ohai.attributes)
+      provides_map = @ohai.attributes
+      provides_map[:zero] = Mash.new
+      provides_map[:zero][:_plugins] = [@plugins[0]]
+      provides_map[:one] = Mash.new
+      provides_map[:one][:_plugins] = [@plugins[1]]
+      provides_map[:one][:two] = Mash.new
+      provides_map[:one][:two][:_plugins] = [@plugins[2]]
+      provides_map[:stub] = Mash.new
+      provides_map[:stub][:three] = Mash.new
+      provides_map[:stub][:three][:_plugins] = [@plugins[3]]
+
+      providers = provides_map.all_plugins
       providers.size.should eql(@plugins.size)
       @plugins.each do |plugin|
         providers.include?(plugin).should be_true
@@ -400,11 +401,11 @@ EOF
         vds['v7plugin'] = @v7plugin
         vds['other'] = @other
 
-        a = @ohai.attributes
-        a[:message] = Mash.new
-        a[:message][:_plugins] = [@v7plugin]
-        a[:other] = Mash.new
-        a[:other][:_plugins] = [@other]
+        dependency_map = @ohai.attributes
+        dependency_map[:message] = Mash.new
+        dependency_map[:message][:_plugins] = [@v7plugin]
+        dependency_map[:other] = Mash.new
+        dependency_map[:other][:_plugins] = [@other]
       end
 
       it "should resolve the v7 plugin dependencies" do