lib/jflow/termination_protector.rb in jflow-0.5.0 vs lib/jflow/termination_protector.rb in jflow-0.5.1

@@ -1,11 +1,17 @@
 require 'net/http'
 require 'json'
-
 module JFlow
   class TerminationProtector
 
+    attr_accessor :asg_client, :ec2_client
+
+    def initialize(asg_client = nil, ec2_client = nil)
+      @asg_client = asg_client || Aws::AutoScaling::Client.new(region: region, credentials: Aws::InstanceProfileCredentials.new)
+      @ec2_client = ec2_client || Aws::EC2::Client.new(region: region, credentials: Aws::InstanceProfileCredentials.new)
+    end
+
     def region
       instance_data['region']
     end
 
     def instance_id
@@ -16,11 +22,10 @@
     def instance_data
       @instance_data ||= JSON.parse(Net::HTTP.get(URI.parse('http://169.254.169.254/latest/dynamic/instance-identity/document')))
     end
 
     def get_asg_name
-      ec2_client = Aws::EC2::Client.new(region: region, credentials: Aws::InstanceProfileCredentials.new)
       instance_tags = ec2_client.describe_tags(filters: [
         {
           name: "resource-id",
           values: [instance_id]
         }
@@ -30,20 +35,23 @@
 
       asg_name
     end
 
     def set_protection(protect_status)
+      @previous_protect_status ||= false
+      return if @previous_protect_status == protect_status
+      @previous_protect_status = protect_status
+
       JFlow.configuration.logger.debug "Setting termination protection status to #{protect_status} for instance #{instance_id} in region #{region}"
       begin
-        asg_client = Aws::AutoScaling::Client.new(region: region, credentials: Aws::InstanceProfileCredentials.new)
         asg_client.set_instance_protection({
           instance_ids: [instance_id],
           auto_scaling_group_name: get_asg_name,
           protected_from_scale_in: protect_status
         })
       rescue => e
         JFlow.configuration.logger.debug "Something went wrong setting termination proection: #{e.inspect}"
+        JFlow.handle_exception(e)
       end
     end
-
   end
 end