From 47bb155125fda78b794a714f717d798631de45d8 Mon Sep 17 00:00:00 2001
From: Daniel Beckham <dbeckham@dealnews.com>
Date: Mon, 14 Dec 2015 16:26:08 -0600
Subject: [PATCH] Add support for tcp check type

---
 .../functions/consul_validate_checks.rb       | 10 +++-
 manifests/check.pp                            |  6 ++
 spec/defines/consul_check_spec.rb             | 58 ++++++++++++++++++-
 spec/functions/consul_validate_checks_spec.rb | 18 ++++++
 4 files changed, 87 insertions(+), 5 deletions(-)

diff --git a/lib/puppet/parser/functions/consul_validate_checks.rb b/lib/puppet/parser/functions/consul_validate_checks.rb
index 26639e23..b08ffed6 100644
--- a/lib/puppet/parser/functions/consul_validate_checks.rb
+++ b/lib/puppet/parser/functions/consul_validate_checks.rb
@@ -15,14 +15,18 @@ def validate_checks(obj)
           end
         elsif obj.key?("http")
           if (! obj.key?("interval"))
-            raise Puppet::ParseError.new('http must be defined for interval checks')
+            raise Puppet::ParseError.new('interval must be defined for http checks')
+          end
+        elsif obj.key?("tcp")
+          if (! obj.key?("interval"))
+            raise Puppet::ParseError.new('interval must be defined for tcp checks')
           end
         elsif obj.key?("script")
           if (! obj.key?("interval"))
-            raise Puppet::ParseError.new('script must be defined for interval checks')
+            raise Puppet::ParseError.new('interval must be defined for script checks')
           end
         else
-          raise Puppet::ParseError.new('One of ttl, script, or http must be defined.')
+          raise Puppet::ParseError.new('One of http, script, tcp, or ttl must be defined.')
         end
     else
       raise Puppet::ParseError.new("Unable to handle object of type <%s>" % obj.class.to_s)
diff --git a/manifests/check.pp b/manifests/check.pp
index 6d1d4139..98fa78b6 100644
--- a/manifests/check.pp
+++ b/manifests/check.pp
@@ -23,6 +23,10 @@
 #   Full path to the location of the healthcheck script. Must be nagios
 #   compliant with regards to the return codes.
 #
+# [*tcp*]
+#   The IP/hostname and port for the service healthcheck. Should be in
+#   'hostname:port' format.
+#
 # [*interval*]
 #   Value in seconds for the interval between runs of the check
 #
@@ -44,6 +48,7 @@
   $ttl        = undef,
   $http       = undef,
   $script     = undef,
+  $tcp        = undef,
   $interval   = undef,
   $service_id = undef,
   $timeout    = undef,
@@ -58,6 +63,7 @@
     'ttl'        => $ttl,
     'http'       => $http,
     'script'     => $script,
+    'tcp'        => $tcp,
     'interval'   => $interval,
     'timeout '   => $timeout,
     'service_id' => $service_id,
diff --git a/spec/defines/consul_check_spec.rb b/spec/defines/consul_check_spec.rb
index a765e296..2e0bf785 100644
--- a/spec/defines/consul_check_spec.rb
+++ b/spec/defines/consul_check_spec.rb
@@ -156,6 +156,52 @@
         .that_notifies("Class[consul::reload_service]") \
     }
   end
+  describe 'with tcp' do
+    let(:params) {{
+      'tcp'      => 'localhost:80',
+      'interval' => '30s',
+    }}
+    it {
+      should contain_file("/etc/consul/check_my_check.json") \
+        .with_content(/"id" *: *"my_check"/) \
+        .with_content(/"name" *: *"my_check"/) \
+        .with_content(/"check" *: *\{/) \
+        .with_content(/"tcp" *: *"localhost:80"/) \
+        .with_content(/"interval" *: *"30s"/)
+    }
+  end
+  describe 'with script and service_id' do
+    let(:params) {{
+      'tcp'        => 'localhost:80',
+      'interval'   => '30s',
+      'service_id' => 'my_service'
+    }}
+    it {
+      should contain_file("/etc/consul/check_my_check.json") \
+        .with_content(/"id" *: *"my_check"/) \
+        .with_content(/"name" *: *"my_check"/) \
+        .with_content(/"check" *: *\{/) \
+        .with_content(/"tcp" *: *"localhost:80"/) \
+        .with_content(/"interval" *: *"30s"/) \
+        .with_content(/"service_id" *: *"my_service"/)
+    }
+  end
+  describe 'reload service with script and token' do
+    let(:params) {{
+      'tcp'      => 'localhost:80',
+      'interval' => '30s',
+      'token'    => 'too-cool-for-this-script'
+    }}
+    it {
+      should contain_file("/etc/consul/check_my_check.json") \
+        .with_content(/"id" *: *"my_check"/) \
+        .with_content(/"name" *: *"my_check"/) \
+        .with_content(/"tcp" *: *"localhost:80"/) \
+        .with_content(/"interval" *: *"30s"/) \
+        .with_content(/"token" *: *"too-cool-for-this-script"/) \
+        .that_notifies("Class[consul::reload_service]") \
+    }
+  end
   describe 'with both ttl and interval' do
     let(:params) {{
       'ttl' => '30s',
@@ -197,7 +243,7 @@
       'script' => 'true',
     }}
     it {
-      should raise_error(Puppet::Error, /script must be defined for interval checks/)
+      should raise_error(Puppet::Error, /interval must be defined for script checks/)
     }
   end
   describe 'with http but no interval' do
@@ -205,7 +251,15 @@
       'http' => 'http://localhost/health',
     }}
     it {
-      should raise_error(Puppet::Error, /http must be defined for interval checks/)
+      should raise_error(Puppet::Error, /interval must be defined for http checks/)
+    }
+  end
+  describe 'with tcp but no interval' do
+    let(:params) {{
+      'tcp' => 'localhost:80',
+    }}
+    it {
+      should raise_error(Puppet::Error, /interval must be defined for tcp checks/)
     }
   end
   describe 'with a / in the id' do
diff --git a/spec/functions/consul_validate_checks_spec.rb b/spec/functions/consul_validate_checks_spec.rb
index c50ae8c2..ee3af8c2 100644
--- a/spec/functions/consul_validate_checks_spec.rb
+++ b/spec/functions/consul_validate_checks_spec.rb
@@ -53,4 +53,22 @@
       }
     ]).and_raise_error(Exception) }
   end
+
+  describe 'validate tcp check' do
+    it {should run.with_params([
+      {
+        'tcp'      => 'localhost:80',
+        'interval' => '30s',
+      }
+    ])}
+  end
+
+  describe 'validate tcp missing interval' do
+    it {should run.with_params([
+      {
+        'tcp' => 'localhost:80'
+      }
+    ]).and_raise_error(Exception) }
+  end
+
 end