From efeb21b9251bc79c65f527bce23c3bcfd4a3e2fc Mon Sep 17 00:00:00 2001
From: Sean Escriva <sean.escriva@gmail.com>
Date: Fri, 17 Mar 2017 17:32:28 -0400
Subject: [ec2] Add additional data from identity document

Since it's already fetching the identity document to retrieve the account id,
this adds the region and AZ info in the same place.

Signed-off-by: Sean Escriva <sean.escriva@gmail.com>
---
 lib/ohai/plugins/ec2.rb       |  2 ++
 spec/unit/plugins/ec2_spec.rb | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/lib/ohai/plugins/ec2.rb b/lib/ohai/plugins/ec2.rb
index 76c718ea..1ad6b88a 100644
--- a/lib/ohai/plugins/ec2.rb
+++ b/lib/ohai/plugins/ec2.rb
@@ -99,6 +99,8 @@ Ohai.plugin(:EC2) do
       end
       ec2[:userdata] = fetch_userdata
       ec2[:account_id] = fetch_dynamic_data["accountId"]
+      ec2[:availability_zone] = fetch_dynamic_data["availabilityZone"]
+      ec2[:region] = fetch_dynamic_data["region"]
       # ASCII-8BIT is equivalent to BINARY in this case
       if ec2[:userdata] && ec2[:userdata].encoding.to_s == "ASCII-8BIT"
         Ohai::Log.debug("Plugin EC2: Binary UserData Found. Storing in base64")
diff --git a/spec/unit/plugins/ec2_spec.rb b/spec/unit/plugins/ec2_spec.rb
index d726ead6..46888c6a 100644
--- a/spec/unit/plugins/ec2_spec.rb
+++ b/spec/unit/plugins/ec2_spec.rb
@@ -127,6 +127,49 @@ describe Ohai::System, "plugin ec2" do
         expect(plugin[:ec2]["account_id"]).to eq("4815162342")
       end
 
+      it "fetches AWS region" do
+        paths.each do |name, body|
+          expect(@http_client).to receive(:get).
+            with("/2012-01-12/#{name}").
+            and_return(double("Net::HTTP Response", :body => body, :code => "200"))
+        end
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/user-data/").
+          and_return(double("Net::HTTP Response", :body => "^_<8B>^H^H<C7>U^@^Csomething^@KT<C8><C9>,)<C9>IU(I-.I<CB><CC>I<E5>^B^@^Qz<BF><B0>^R^@^@^@", :code => "200"))
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/dynamic/instance-identity/document/").
+          and_return(double("Net::HTTP Response", :body => "{\"region\":\"us-east-1\"}", :code => "200"))
+
+        plugin.run
+
+        expect(plugin[:ec2]).not_to be_nil
+        expect(plugin[:ec2]["instance_type"]).to eq("c1.medium")
+        expect(plugin[:ec2]["ami_id"]).to eq("ami-5d2dc934")
+        expect(plugin[:ec2]["security_groups"]).to eql %w{group1 group2}
+        expect(plugin[:ec2]["region"]).to eq("us-east-1")
+      end
+
+      it "fetches AWS availability zone" do
+        paths.each do |name, body|
+          expect(@http_client).to receive(:get).
+            with("/2012-01-12/#{name}").
+            and_return(double("Net::HTTP Response", :body => body, :code => "200"))
+        end
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/user-data/").
+          and_return(double("Net::HTTP Response", :body => "^_<8B>^H^H<C7>U^@^Csomething^@KT<C8><C9>,)<C9>IU(I-.I<CB><CC>I<E5>^B^@^Qz<BF><B0>^R^@^@^@", :code => "200"))
+        expect(@http_client).to receive(:get).
+          with("/2012-01-12/dynamic/instance-identity/document/").
+          and_return(double("Net::HTTP Response", :body => "{\"availabilityZone\":\"us-east-1d\"}", :code => "200"))
+
+        plugin.run
+
+        expect(plugin[:ec2]).not_to be_nil
+        expect(plugin[:ec2]["instance_type"]).to eq("c1.medium")
+        expect(plugin[:ec2]["ami_id"]).to eq("ami-5d2dc934")
+        expect(plugin[:ec2]["security_groups"]).to eql %w{group1 group2}
+        expect(plugin[:ec2]["availability_zone"]).to eq("us-east-1d")
+      end
     end
 
     it "parses ec2 network/ directory as a multi-level hash" do
-- 
cgit v1.2.1