diff --git a/lib/scientist/experiment.rb b/lib/scientist/experiment.rb index f48b339..5720ac4 100644 --- a/lib/scientist/experiment.rb +++ b/lib/scientist/experiment.rb @@ -290,6 +290,13 @@ def use(&block) try "control", &block end + # Define a block which will determine the cohort of this experiment + # when called. The block will be passed a `Scientist::Result` as its + # only argument and the cohort will be set on the result. + def cohort(&block) + @_scientist_determine_cohort = block + end + # Whether or not to raise a mismatch error when a mismatch occurs. def raise_on_mismatches? if raise_on_mismatches.nil? @@ -316,7 +323,7 @@ def generate_result(name) end control = observations.detect { |o| o.name == name } - Scientist::Result.new(self, observations, control) + Scientist::Result.new(self, observations, control, @_scientist_determine_cohort) end private diff --git a/lib/scientist/result.rb b/lib/scientist/result.rb index 76a4d21..79dc43b 100644 --- a/lib/scientist/result.rb +++ b/lib/scientist/result.rb @@ -19,19 +19,33 @@ class Scientist::Result # An Array of Observations in execution order. attr_reader :observations + # If the experiment was defined with a cohort block, the cohort this + # result has been determined to belong to. + attr_reader :cohort + # Internal: Create a new result. # - # experiment - the Experiment this result is for - # observations: - an Array of Observations, in execution order - # control: - the control Observation + # experiment - the Experiment this result is for + # observations: - an Array of Observations, in execution order + # control: - the control Observation + # determine_cohort - An optional callable that is passed the Result to + # determine its cohort # - def initialize(experiment, observations = [], control = nil) + def initialize(experiment, observations = [], control = nil, determine_cohort = nil) @experiment = experiment @observations = observations @control = control @candidates = observations - [control] evaluate_candidates + if determine_cohort + begin + @cohort = determine_cohort.call(self) + rescue StandardError => e + experiment.raised :cohort, e + end + end + freeze end diff --git a/test/scientist/experiment_test.rb b/test/scientist/experiment_test.rb index 12462f6..87240e8 100644 --- a/test/scientist/experiment_test.rb +++ b/test/scientist/experiment_test.rb @@ -302,6 +302,46 @@ def @ex.enabled? assert_equal "kaboom", exception.message end + describe "cohorts" do + it "accepts a cohort config block" do + @ex.cohort { "1" } + end + + it "assigns a cohort to the result using the provided block" do + @ex.context(foo: "bar") + @ex.cohort { |res| "foo-#{res.context[:foo]}-#{Math.log10(res.control.value).round}" } + @ex.use { 5670 } + @ex.try { 5670 } + + @ex.run + assert_equal "foo-bar-4", @ex.published_result.cohort + end + + it "assigns no cohort if no cohort block passed" do + @ex.use { 5670 } + @ex.try { 5670 } + + @ex.run + assert_nil @ex.published_result.cohort + end + + it "rescues errors raised in the cohort determination block" do + @ex.use { 5670 } + @ex.try { 5670 } + @ex.cohort { |res| raise "intentional" } + + @ex.run + + refute_nil @ex.published_result + assert_nil @ex.published_result.cohort + + assert_equal 1, @ex.exceptions.size + code, exception = @ex.exceptions[0] + assert_equal :cohort, code + assert_equal "intentional", exception.message + end + end + describe "#raise_with" do it "raises custom error if provided" do CustomError = Class.new(Scientist::Experiment::MismatchError) diff --git a/test/scientist/result_test.rb b/test/scientist/result_test.rb index c9bc41e..474a851 100644 --- a/test/scientist/result_test.rb +++ b/test/scientist/result_test.rb @@ -98,6 +98,17 @@ assert_equal @experiment.name, result.experiment_name end + it "takes an optional callable to determine cohort" do + a = Scientist::Observation.new("a", @experiment) { 1 } + b = Scientist::Observation.new("b", @experiment) { 1 } + + result = Scientist::Result.new @experiment, [a, b], a + assert_nil result.cohort + + result = Scientist::Result.new @experiment, [a, b], a, ->(res) { "cohort-1" } + assert_equal "cohort-1", result.cohort + end + it "has the context from an experiment" do @experiment.context :foo => :bar a = Scientist::Observation.new("a", @experiment) { 1 }