From 8f02c34360c203dd91a4ce981b249aaaa60458e6 Mon Sep 17 00:00:00 2001
From: Pawel Pikula <pawel.pikula@pattern-match.com>
Date: Thu, 28 May 2020 00:18:27 +0200
Subject: [PATCH] support private DNS names if public one is not available

---
 lib/elbas/aws/instance.rb            |  9 +++++++--
 lib/elbas/aws/instance_collection.rb |  2 +-
 spec/aws/ami_spec.rb                 |  4 ++--
 spec/aws/instance_spec.rb            | 25 +++++++++++++++----------
 4 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/lib/elbas/aws/instance.rb b/lib/elbas/aws/instance.rb
index 0530215..ad8e5e8 100644
--- a/lib/elbas/aws/instance.rb
+++ b/lib/elbas/aws/instance.rb
@@ -5,15 +5,20 @@ class Instance < Base
 
       attr_reader :aws_counterpart, :id, :state
 
-      def initialize(id, public_dns, state)
+      def initialize(id, public_dns, private_dns, state)
         @id = id
         @public_dns = public_dns
+        @private_dns = private_dns
         @state = state
         @aws_counterpart = aws_namespace::Instance.new id, client: aws_client
       end
 
       def hostname
-        @public_dns
+          if @public_dns.nil? || @public_dns.empty?
+              @private_dns
+          else
+              @public_dns
+          end
       end
 
       def running?
diff --git a/lib/elbas/aws/instance_collection.rb b/lib/elbas/aws/instance_collection.rb
index c0ac6cf..970235e 100644
--- a/lib/elbas/aws/instance_collection.rb
+++ b/lib/elbas/aws/instance_collection.rb
@@ -8,7 +8,7 @@ class InstanceCollection < Base
       def initialize(ids)
         @ids = ids
         @instances = query_instances_by_ids(ids).map do |i|
-          Instance.new(i.instance_id, i.public_dns_name, i.state.code)
+          Instance.new(i.instance_id, i.public_dns_name, i.private_dns_name, i.state.code)
         end
       end
 
diff --git a/spec/aws/ami_spec.rb b/spec/aws/ami_spec.rb
index f75120e..a00a71d 100644
--- a/spec/aws/ami_spec.rb
+++ b/spec/aws/ami_spec.rb
@@ -104,7 +104,7 @@
 
   describe '.create' do
     subject { described_class }
-    let(:instance) { Elbas::AWS::Instance.new 'i-1234567890', nil, nil }
+    let(:instance) { Elbas::AWS::Instance.new 'i-1234567890', nil, nil, nil }
 
     before do
       webmock :post, %r{ec2.(.*).amazonaws.com\/\z} => 'CreateImage.200.xml',
@@ -137,4 +137,4 @@
       expect(ami.id).to eq 'ami-4fa54026' # from stub
     end
   end
-end
\ No newline at end of file
+end
diff --git a/spec/aws/instance_spec.rb b/spec/aws/instance_spec.rb
index def2643..599ce4e 100644
--- a/spec/aws/instance_spec.rb
+++ b/spec/aws/instance_spec.rb
@@ -1,30 +1,35 @@
 describe Elbas::AWS::Instance do
-  subject { Elbas::AWS::Instance.new 'i-1234567890', 'ec2-1234567890.amazonaws.com', 16 }
+  let (:public_instance) { Elbas::AWS::Instance.new 'i-1234567890', 'ec2-1234567890.amazonaws.com', 'ec2-1234567890.internal',  16 }
+  let (:private_instance) { Elbas::AWS::Instance.new 'i-1234567890', nil, 'ec2-1234567890.internal',  16 }
 
   describe '#initialize' do
     it 'sets the AWS counterpart' do
-      expect(subject.aws_counterpart).to be_a_kind_of ::Aws::EC2::Instance
-      expect(subject.aws_counterpart.id).to eq 'i-1234567890'
+      expect(public_instance.aws_counterpart).to be_a_kind_of ::Aws::EC2::Instance
+      expect(public_instance.aws_counterpart.id).to eq 'i-1234567890'
     end
   end
 
   describe '#hostname' do
-    it 'returns the public DNS' do
-      expect(subject.hostname).to eq 'ec2-1234567890.amazonaws.com'
+    it 'returns the public DNS if it is set' do
+      expect(public_instance.hostname).to eq 'ec2-1234567890.amazonaws.com'
+    end
+
+    it 'returns private DNS if public one is not available' do
+        expect(private_instance.hostname).to eq 'ec2-1234567890.internal'
     end
   end
 
   describe '#running?' do
     it 'returns true if the state code is 16 ("running")' do
-      expect(subject).to receive(:state) { 16 }
-      expect(subject).to be_running
+      expect(public_instance).to receive(:state) { 16 }
+      expect(public_instance).to be_running
     end
 
     it 'returns false for every other state code' do
       [0, 32, 48, 64, 80].each do |code|
-        expect(subject).to receive(:state) { code }
-        expect(subject).to_not be_running
+        expect(public_instance).to receive(:state) { code }
+        expect(public_instance).to_not be_running
       end
     end
   end
-end
\ No newline at end of file
+end