Extend set_attribute to set sub-attributes.OHAI-794/set-attribute

Use:
  - set_attribute(:attr, value) assigns attribute 'attr' value.
  - set_attribute(:attr, :subattr, value) assigns attribute 'attr/subattr' value. Extends to subattributes of arbitrary depth.

Raises a TypeError if some :subattr is not an Array, Hash, or Mash (e.g., a String value).

2 files changed, 71 insertions, 9 deletions

diff --git a/lib/ohai/dsl/plugin.rb b/lib/ohai/dsl/plugin.rb
index adcb6f82..b9fe9910 100644
--- a/lib/ohai/dsl/plugin.rb
+++ b/lib/ohai/dsl/plugin.rb
@@ -152,8 +152,19 @@ module Ohai
         end
       end
 
-      def set_attribute(name, *values)
-        @data[name] = Array18(*values)
+      def set_attribute(name, *attrs, value)
+        # Initialize the path in the @data Mash with new Mashes, if needed.
+        # Will raise a TypeError if we hit a subattribute that is not a
+        # Hash, Mash, or Array.
+        keys = [name] + attrs
+        attribute = keys[0..-2].inject(@data) do |attrs, key|
+          attrs[key] ||= Mash.new
+          attrs[key]
+        end
+
+        # Set the subattribute to the value.
+        attr_name = attrs.empty? ? name : attrs[-1]
+        attribute[attr_name] = value
         @data[name]
       end
 
@@ -196,12 +207,6 @@ module Ohai
         # NoMethodError occurs when trying to access a key on nil
         nil
       end
-
-      def Array18(*args)
-        return nil if args.empty?
-        return args.first if args.length == 1
-        return *args
-      end
     end
   end
 end
diff --git a/spec/unit/dsl/plugin_spec.rb b/spec/unit/dsl/plugin_spec.rb
index 2a35701c..c37b6fd2 100644
--- a/spec/unit/dsl/plugin_spec.rb
+++ b/spec/unit/dsl/plugin_spec.rb
@@ -90,11 +90,68 @@ shared_examples "Ohai::DSL::Plugin" do
     end
   end
 
-  context "when setting attributes" do
+  describe "set_attribute" do
+    it "raises an ArgumentError when given one argument" do
+      expect { plugin.set_attribute(:tea) }.to raise_error(ArgumentError)
+    end
+
     it "lets you set an attribute" do
       plugin.set_attribute(:tea, "is soothing")
       expect(plugin.data["tea"]).to eql("is soothing")
     end
+
+    it "lets you set an attribute to an empty array" do
+      plugin.set_attribute(:tea, [])
+      expect(plugin.data["tea"]).to eql([])
+    end
+
+    it "lets you set an attribute to an empty Mash" do
+      plugin.set_attribute(:tea, {})
+      expect(plugin.data["tea"]).to eql({})
+    end
+
+    it "lets you modify an attribute" do
+      plugin.data["tea"] = "is soothing"
+      plugin.set_attribute(:tea, "tastes like plants")
+      expect(plugin.data["tea"]).to eql("tastes like plants")
+    end
+
+    it "lets you set a subattribute" do
+      plugin.data["tea"] = Mash.new
+      plugin.set_attribute(:tea, :green, "tastes like green")
+      expect(plugin.data["tea"]["green"]).to eql("tastes like green")
+    end
+
+    it "lets you set a subattribute to an empty array" do
+      plugin.set_attribute(:tea, :green, [])
+      expect(plugin.data["tea"]["green"]).to eql([])
+    end
+
+    it "lets you set a subattribute to an empty Mash" do
+      plugin.set_attribute(:tea, :green, {})
+      expect(plugin.data["tea"]["green"]).to eql({})
+    end
+
+    it "lets you modify a subattribute" do
+      plugin.data["tea"] = Mash.new
+      plugin.data["tea"]["green"] = "tastes like green"
+      plugin.set_attribute(:tea, :green, "made from leaves")
+      expect(plugin.data["tea"]["green"]).to eql("made from leaves")
+    end
+
+    it "raises a TypeError when modifying an attribute that is not a Mash" do
+      plugin.data["tea"] = Mash.new
+      plugin.data["tea"] = "is soothing"
+      expect { plugin.set_attribute(:tea, :green, "tastes like green") }.to raise_error(TypeError)
+    end
+
+    it "lets you modify a subattribute with an array" do
+      green = { flavor: "tastes like green" }
+      black = { flavor: "tastes like black" }
+      plugin.data[:tea] = [green, black]
+      plugin.set_attribute(:tea, 0, :source, "made from leaves")
+      expect(plugin.data[:tea][0][:source]).to eq("made from leaves")
+    end
   end
 
   context "when getting attributes" do