bin/check-marathon-task.rb in sensu-plugins-mesos-2.0.0 vs bin/check-marathon-task.rb in sensu-plugins-mesos-2.1.0

@@ -28,17 +28,17 @@
 #   Released under the same terms as Sensu (the MIT license); see LICENSE
 #   for details.
 #
 
 require 'sensu-plugin/check/cli'
-require 'net/http'
+require 'rest-client'
 require 'json'
 
 # This plugin checks that the given Mesos/Marathon task is running properly.
 #
 # This means that all of the following is true:
-# 1. There are N tasks for the app, as defined by the --instances parameter
+# 1. There are N tasks for the app, as defined by the --instances parameter or checks configured tasks in Marathon as fallback
 # 2. Each task's state is running
 # 3. No task is unhealthy, as defined in Marathon
 #
 # A task is seen as **unhealthy** by Marathon if any of the health checks for
 # the task is not **alive**.  Alive means that a check has a last success that
@@ -69,11 +69,12 @@
          required: true
 
   option :instances,
          short: '-i INSTANCES',
          long: '--instances INSTANCES',
-         required: true,
+         required: false,
+         default: 0,
          proc: proc(&:to_i)
 
   option :protocol,
          short: '-P PROTOCOL',
          long: '--protocol PROTOCOL',
@@ -87,47 +88,43 @@
 
   option :password,
          long: '--password PASSWORD',
          required: false
 
-  def run
-    if config[:instances].zero?
-      unknown 'number of instances should be an integer'
-    end
+  option :timeout,
+         description: 'timeout in seconds',
+         short: '-T TIMEOUT',
+         long: '--timeout TIMEOUT',
+         proc: proc(&:to_i),
+         default: 5
 
+  def run
     if !config[:username].nil? && config[:password].nil? ||
        config[:username].nil? && !config[:password].nil?
       unknown 'You must provide both username and password'
     end
 
     failures = []
     uri = config[:uri]
     config[:server].split(',').each do |s|
       begin
-        url = URI.parse("#{config[:protocol]}://#{s}:#{config[:port]}#{uri}")
-        req = Net::HTTP::Get.new(url)
-        req.add_field('Accept', 'application/json')
-        if !config[:username].nil? && !config[:password].nil?
-          req.basic_auth(config[:username], config[:password])
-        end
-        r = Net::HTTP.start(url.host, url.port,
-                            use_ssl: config[:protocol] == 'https') do |h|
-          h.request(req)
-        end
+        auth_headers = {}
+        auth_headers = { Authorization: "#{config[:username]} #{config[:password]}" } if !config[:username].nil? && !config[:password].nil?
+        r = RestClient::Resource.new("#{config[:protocol]}://#{s}:#{config[:port]}#{uri}", auth_headers, config[:timeout]).get
+        expected = if config[:instances].zero?
+                     default_tasks(s)
+                   else
+                     config[:instances]
+                   end
+        ok_count, unhealthy = check_tasks r
 
-        ok_count, unhealthy = check_tasks r.body
+        message = "#{ok_count}/#{expected} #{config[:task]} tasks running"
 
-        message = "#{ok_count}/#{config[:instances]} #{config[:task]} tasks running"
+        message << ":\n" << unhealthy.join("\n") if unhealthy.any?
 
-        if unhealthy.any?
-          message << ":\n" << unhealthy.join("\n")
-        end
+        critical message if unhealthy.any? || ok_count < config[:instances]
 
-        if unhealthy.any? || ok_count < config[:instances]
-          critical message
-        end
-
         ok message
       rescue Errno::ECONNREFUSED, SocketError
         failures << "Marathon on #{s} could not be reached"
       rescue => err
         failures << "error caught trying to reach Marathon on #{s}: #{err}"
@@ -170,7 +167,16 @@
         unhealthy << message
       end
     end
 
     [tasks.length, unhealthy]
+  end
+
+  def default_tasks(server)
+    expected_tasks_url = "/v2/apps/#{config[:task]}"
+    auth_headers = {}
+    auth_headers = { Authorization: "#{config[:username]} #{config[:password]}" } if !config[:username].nil? && !config[:password].nil?
+    r = RestClient::Resource.new("#{config[:protocol]}://#{server}:#{config[:port]}#{expected_tasks_url}", auth_headers, config[:timeout]).get
+    n_tasks = JSON.parse(r)['app']['instances']
+    n_tasks
   end
 end