From 9d8c53cf31e0f1ede0f6ed291ca4bb4825396197 Mon Sep 17 00:00:00 2001
From: Adam Leff <adam@leff.co>
Date: Fri, 6 Oct 2017 13:24:31 -0400
Subject: [PATCH] Support symbol keys in ObjectTraverser (#2221)

As detected in #2036, it is not possible to extract values from
a YAML file if the key is a symbol. This change refactors ObjectTraverser
to support symbol keys before attempting to stringify them.

Signed-off-by: Adam Leff <adam@leff.co>
---
 lib/utils/object_traversal.rb            | 47 ++++++++++++++++++------
 test/unit/mock/files/kitchen.yml         |  3 ++
 test/unit/resources/yaml_test.rb         |  8 ++++
 test/unit/utils/object_traversal_test.rb | 13 ++++++-
 4 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/lib/utils/object_traversal.rb b/lib/utils/object_traversal.rb
index b10574c93..5d760a026 100644
--- a/lib/utils/object_traversal.rb
+++ b/lib/utils/object_traversal.rb
@@ -3,22 +3,47 @@
 # author: Christoph Hartmann
 module ObjectTraverser
   def extract_value(keys, value)
-    key = keys.shift
-    return nil if key.nil? || value.nil?
+    return nil if value.nil?
 
-    if value.is_a?(Array)
-      value = if key.is_a?(Fixnum)
-                value[key]
-              elsif value.respond_to?(key.to_sym)
-                value.send(key.to_sym)
-              end
-    else
-      value = value[key.to_s].nil? ? nil : value[key.to_s]
-    end
+    key = keys.shift
+    return nil if key.nil?
+
+    # if the current value is not a Hash or Array, it is undefined
+    # behavior so value will be assigned nil by default.
+    value = if value.is_a?(Array)
+              extract_from_array(key, value)
+            elsif value.is_a?(Hash)
+              extract_from_hash(key, value)
+            end
 
     # if there are no more keys, just return the value
     return value if keys.first.nil?
     # if there are more keys, extract more
     extract_value(keys.clone, value)
   end
+
+  private
+
+  # If the values to return from is an Array, allow returning by index.
+  # Otherwise, support methods on the Array itself.
+  def extract_from_array(key, value)
+    if key.is_a?(Fixnum)
+      value[key]
+    elsif value.respond_to?(key.to_sym)
+      value.send(key.to_sym)
+    end
+  end
+
+  # for Hashes, try to return the value by the key.
+  # We first try to find by the raw key before we stringify
+  # if the keys themselves are symbols, for example.
+  #
+  # This will return nil default if we can't find the key.
+  def extract_from_hash(key, value)
+    if value.key?(key)
+      value[key]
+    elsif value.key?(key.to_s)
+      value[key.to_s]
+    end
+  end
 end
diff --git a/test/unit/mock/files/kitchen.yml b/test/unit/mock/files/kitchen.yml
index 2bf751297..8cca89304 100644
--- a/test/unit/mock/files/kitchen.yml
+++ b/test/unit/mock/files/kitchen.yml
@@ -5,3 +5,6 @@ driver:
 platforms:
   - linux
   - mac
+:symbol_key: 123
+:symbol_key_deep:
+  foo: bar
diff --git a/test/unit/resources/yaml_test.rb b/test/unit/resources/yaml_test.rb
index 6b8974b24..d2269531b 100644
--- a/test/unit/resources/yaml_test.rb
+++ b/test/unit/resources/yaml_test.rb
@@ -30,5 +30,13 @@ describe 'Inspec::Resources::YAML' do
     it 'doesnt resolve symbol-notation names' do
       _(resource.send(:'driver.customize.memory')).must_be_nil
     end
+
+    it 'supports fetching by symbol keys' do
+      _(resource.send(:symbol_key)).must_equal 123
+    end
+
+    it 'support fetching by symbol keys in array syntax for rspec-its' do
+      _(resource.send(:[], :symbol_key_deep, 'foo')).must_equal 'bar'
+    end
   end
 end
diff --git a/test/unit/utils/object_traversal_test.rb b/test/unit/utils/object_traversal_test.rb
index 00558e287..dd4520bd2 100644
--- a/test/unit/utils/object_traversal_test.rb
+++ b/test/unit/utils/object_traversal_test.rb
@@ -25,7 +25,12 @@ describe ObjectTraverser do
         123,
         456,
         { 'array1hashkey1' => 1, 'array1hashkey2' => 2 },
-      ]
+      ],
+      :symbol_key_1 => 123,
+      :symbol_key_2 => {
+        :symbol_under_symbol => 456,
+        'string_under_symbol' => 789
+      }
     }
   end
 
@@ -77,4 +82,10 @@ describe ObjectTraverser do
     subject.extract_value(['array2', 'last', 'array1hashkey1'], sample_data).must_equal(1)
     subject.extract_value(['array2', 'last', 'array1hashkey2'], sample_data).must_equal(2)
   end
+
+  it 'supports returning values with symbol keys' do
+    subject.extract_value([:symbol_key_1], sample_data).must_equal(123)
+    subject.extract_value([:symbol_key_2, :symbol_under_symbol], sample_data).must_equal(456)
+    subject.extract_value([:symbol_key_2, 'string_under_symbol'], sample_data).must_equal(789)
+  end
 end