From f4c6220bad5d92e1dcb15ea8757565a0b63b1d63 Mon Sep 17 00:00:00 2001
From: Trey Dockendorf <treydock@gmail.com>
Date: Fri, 22 Jan 2021 12:24:56 -0500
Subject: [PATCH] Add scheduler property to sensu_check

---
 lib/puppet/type/sensu_check.rb      |  5 +++++
 spec/acceptance/sensu_check_spec.rb |  5 +++++
 spec/unit/sensu_check_spec.rb       | 22 ++++++++++++++++++++++
 3 files changed, 32 insertions(+)

diff --git a/lib/puppet/type/sensu_check.rb b/lib/puppet/type/sensu_check.rb
index 8810f84d3b..db2c69f061 100644
--- a/lib/puppet/type/sensu_check.rb
+++ b/lib/puppet/type/sensu_check.rb
@@ -229,6 +229,11 @@
     newvalues(/.*/, :absent)
   end
 
+  newproperty(:scheduler) do
+    desc "Type of scheduler that schedules the check."
+    newvalues(:memory, :etcd, :postgres)
+  end
+
   newproperty(:output_metric_format) do
     desc "The metric format generated by the check command."
     newvalues(:nagios_perfdata, :graphite_plaintext, :influxdb_line, :opentsdb_line, :prometheus_text, :absent)
diff --git a/spec/acceptance/sensu_check_spec.rb b/spec/acceptance/sensu_check_spec.rb
index 37d06d6983..5cbc419d95 100644
--- a/spec/acceptance/sensu_check_spec.rb
+++ b/spec/acceptance/sensu_check_spec.rb
@@ -19,6 +19,7 @@
         proxy_requests                   => {
           'entity_attributes' => ["entity.Class == 'proxy'"],
         },
+        scheduler                        => 'memory',
         output_metric_format             => 'nagios_perfdata',
         output_metric_tags               => [{'name' => 'instance', 'value' => '{{ .name }}'}],
         labels                           => { 'foo' => 'baz' },
@@ -76,6 +77,7 @@
         expect(data['stdin']).to eq(false)
         expect(data['check_hooks']).to eq([{'0' => ['always.sh']},{'1' => ['test.sh']},{'critical' => ['httpd-restart']}])
         expect(data['proxy_requests']['entity_attributes']).to eq(["entity.Class == 'proxy'"])
+        expect(data['scheduler']).to eq('memory')
         expect(data['output_metric_format']).to eq('nagios_perfdata')
         expect(data['output_metric_tags']).to eq([{'name' => 'instance', 'value' => '{{ .name }}'}])
         expect(data['metadata']['labels']['foo']).to eq('baz')
@@ -162,6 +164,8 @@ class { 'sensu::cli':
         proxy_requests                   => {
           'entity_attributes' => ['System.OS==linux'],
         },
+        scheduler                        => 'etcd',
+        round_robin                      => true,
         output_metric_format             => 'graphite_plaintext',
         output_metric_tags               => [
           {'name' => 'instance', 'value' => '{{ .name }}'},
@@ -206,6 +210,7 @@ class { 'sensu::cli':
         data = JSON.parse(stdout)
         expect(data['check_hooks']).to eq([{'critical' => ['httpd-restart']},{'warning' => ['httpd-restart']}])
         expect(data['proxy_requests']['entity_attributes']).to eq(['System.OS==linux'])
+        expect(data['scheduler']).to eq('etcd')
         expect(data['output_metric_format']).to eq('graphite_plaintext')
         expect(data['output_metric_tags']).to include({'name' => 'instance', 'value' => '{{ .name }}'})
         expect(data['output_metric_tags']).to include({'name' => 'prometheus_type', 'value' => 'gauge'})
diff --git a/spec/unit/sensu_check_spec.rb b/spec/unit/sensu_check_spec.rb
index d0001972f4..b190f2aac3 100644
--- a/spec/unit/sensu_check_spec.rb
+++ b/spec/unit/sensu_check_spec.rb
@@ -311,6 +311,28 @@
     end
   end
 
+  describe 'scheduler' do
+    [
+      'memory',
+      'etcd',
+      'postgres',
+    ].each do |v|
+      it "should accept #{v}" do
+        config[:scheduler] = v
+        expect(check[:scheduler]).to eq(v.to_sym)
+      end
+    end
+
+    it 'should not have a default' do
+      expect(check[:scheduler]).to be_nil
+    end
+
+    it 'should not accept invalid values' do
+      config[:scheduler] = 'foo'
+      expect { check }.to raise_error(Puppet::Error, /Invalid value "foo"/)
+    end
+  end
+
   describe 'output_metric_format' do
     [
       'nagios_perfdata',