(#2886) seeded_rand: new function

seeded_rand is needed for repeatable randomness across nodes in a cluster

3 files changed, 81 insertions, 0 deletions

diff --git a/README.markdown b/README.markdown
index 526b573..9348a3a 100644
--- a/README.markdown
+++ b/README.markdown
@@ -697,6 +697,12 @@ Reverses the order of a string or array. *Type*: rvalue.
 
 Strips spaces to the right of the string. *Type*: rvalue.
 
+#### `seeded_rand`
+
+Takes an integer max value and a string seed value and returns a
+repeatable random integer smaller than max.  Like `fqdn_rand`, but
+this does not add node specific data to the seed.  *Type*: rvalue.
+
 #### `shuffle`
 
 Randomizes the order of a string or array elements. *Type*: rvalue.
diff --git a/lib/puppet/parser/functions/seeded_rand.rb b/lib/puppet/parser/functions/seeded_rand.rb
new file mode 100644
index 0000000..44e27b8
--- /dev/null
+++ b/lib/puppet/parser/functions/seeded_rand.rb
@@ -0,0 +1,22 @@
+Puppet::Parser::Functions::newfunction(
+  :seeded_rand,
+  :arity => 2,
+  :type => :rvalue,
+  :doc => <<-EOS
+Usage: `seeded_rand(MAX, SEED)`. MAX must be a positive integer; SEED is any string.
+
+Generates a random whole number greater than or equal to 0 and less
+than MAX, using the value of SEED for repeatable randomness.  If SEED
+starts with "$fqdn:", this is behaves the same as `fqdn_rand`.
+
+EOS
+) do |args|
+  require 'digest/md5'
+
+  raise(ArgumentError, "seeded_rand(): first argument must be a positive integer") unless function_is_integer([args[0]]) and args[0].to_i > 0
+  raise(ArgumentError, "seeded_rand(): second argument must be a string") unless args[1].is_a? String
+
+  max = args[0].to_i
+  seed = Digest::MD5.hexdigest(args[1]).hex
+  Puppet::Util.deterministic_rand(seed,max)
+end
diff --git a/spec/functions/seeded_rand_spec.rb b/spec/functions/seeded_rand_spec.rb
new file mode 100644
index 0000000..38e134e
--- /dev/null
+++ b/spec/functions/seeded_rand_spec.rb
@@ -0,0 +1,53 @@
+require 'spec_helper'
+
+describe 'seeded_rand' do
+  it { is_expected.not_to eq(nil) }
+  it { is_expected.to run.with_params().and_raise_error(ArgumentError, /wrong number of arguments/i) }
+  it { is_expected.to run.with_params(1).and_raise_error(ArgumentError, /wrong number of arguments/i) }
+  it { is_expected.to run.with_params(0, '').and_raise_error(ArgumentError, /first argument must be a positive integer/) }
+  it { is_expected.to run.with_params(1.5, '').and_raise_error(ArgumentError, /first argument must be a positive integer/) }
+  it { is_expected.to run.with_params(-10, '').and_raise_error(ArgumentError, /first argument must be a positive integer/) }
+  it { is_expected.to run.with_params("-10", '').and_raise_error(ArgumentError, /first argument must be a positive integer/) }
+  it { is_expected.to run.with_params("string", '').and_raise_error(ArgumentError, /first argument must be a positive integer/) }
+  it { is_expected.to run.with_params([], '').and_raise_error(ArgumentError, /first argument must be a positive integer/) }
+  it { is_expected.to run.with_params({}, '').and_raise_error(ArgumentError, /first argument must be a positive integer/) }
+  it { is_expected.to run.with_params(1, 1).and_raise_error(ArgumentError, /second argument must be a string/) }
+  it { is_expected.to run.with_params(1, []).and_raise_error(ArgumentError, /second argument must be a string/) }
+  it { is_expected.to run.with_params(1, {}).and_raise_error(ArgumentError, /second argument must be a string/) }
+
+  it "provides a random number strictly less than the given max" do
+    expect(seeded_rand(3, 'seed')).to satisfy {|n| n.to_i < 3 }
+  end
+
+  it "provides a random number greater or equal to zero" do
+    expect(seeded_rand(3, 'seed')).to satisfy {|n| n.to_i >= 0 }
+  end
+
+  it "provides the same 'random' value on subsequent calls for the same host" do
+    expect(seeded_rand(10, 'seed')).to eql(seeded_rand(10, 'seed'))
+  end
+
+  it "allows seed to control the random value on a single host" do
+    first_random = seeded_rand(1000, 'seed1')
+    second_different_random = seeded_rand(1000, 'seed2')
+
+    expect(first_random).not_to eql(second_different_random)
+  end
+
+  it "should not return different values for different hosts" do
+    val1 = seeded_rand(1000, 'foo', :host => "first.host.com")
+    val2 = seeded_rand(1000, 'foo', :host => "second.host.com")
+
+    expect(val1).to eql(val2)
+  end
+
+  def seeded_rand(max, seed, args = {})
+    host = args[:host] || '127.0.0.1'
+
+    # workaround not being able to use let(:facts) because some tests need
+    # multiple different hostnames in one context
+    scope.stubs(:lookupvar).with("::fqdn", {}).returns(host)
+
+    scope.function_seeded_rand([max, seed])
+  end
+end