1 files changed, 89 insertions, 5 deletions

diff --git a/spec/lib/gitlab/experiment/rollout/feature_spec.rb b/spec/lib/gitlab/experiment/rollout/feature_spec.rb
index cd46e7b3386..6d01b7a175f 100644
--- a/spec/lib/gitlab/experiment/rollout/feature_spec.rb
+++ b/spec/lib/gitlab/experiment/rollout/feature_spec.rb
@@ -2,16 +2,15 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Experiment::Rollout::Feature, :experiment do
+RSpec.describe Gitlab::Experiment::Rollout::Feature, :experiment, feature_category: :acquisition do
   subject { described_class.new(subject_experiment) }
 
   let(:subject_experiment) { experiment('namespaced/stub') }
 
-  describe "#enabled?" do
+  describe "#enabled?", :saas do
     before do
       stub_feature_flags(gitlab_experiment: true)
       allow(subject).to receive(:feature_flag_defined?).and_return(true)
-      allow(Gitlab).to receive(:com?).and_return(true)
       allow(subject).to receive(:feature_flag_instance).and_return(double(state: :on))
     end
 
@@ -45,6 +44,18 @@ RSpec.describe Gitlab::Experiment::Rollout::Feature, :experiment do
   end
 
   describe "#execute_assignment" do
+    let(:variants) do
+      ->(e) do
+        # rubocop:disable Lint/EmptyBlock
+        e.control {}
+        e.variant(:red) {}
+        e.variant(:blue) {}
+        # rubocop:enable Lint/EmptyBlock
+      end
+    end
+
+    let(:subject_experiment) { experiment('namespaced/stub', &variants) }
+
     before do
       allow(Feature).to receive(:enabled?).with('namespaced_stub', any_args).and_return(true)
     end
@@ -60,9 +71,82 @@ RSpec.describe Gitlab::Experiment::Rollout::Feature, :experiment do
     end
 
     it "returns an assigned name" do
-      allow(subject).to receive(:behavior_names).and_return([:variant1, :variant2])
+      expect(subject.execute_assignment).to eq(:blue)
+    end
+
+    context "when there are no behaviors" do
+      let(:variants) { ->(e) { e.control {} } } # rubocop:disable Lint/EmptyBlock
+
+      it "does not raise an error" do
+        expect { subject.execute_assignment }.not_to raise_error
+      end
+    end
+
+    context "for even rollout to non-control", :saas do
+      let(:counts) { Hash.new(0) }
+      let(:subject_experiment) { experiment('namespaced/stub') }
+
+      before do
+        allow_next_instance_of(described_class) do |instance|
+          allow(instance).to receive(:enabled?).and_return(true)
+        end
+
+        subject_experiment.variant(:variant1) {} # rubocop:disable Lint/EmptyBlock
+        subject_experiment.variant(:variant2) {} # rubocop:disable Lint/EmptyBlock
+      end
+
+      it "rolls out relatively evenly to 2 behaviors" do
+        100.times { |i| run_cycle(subject_experiment, value: i) }
+
+        expect(counts).to eq(variant1: 54, variant2: 46)
+      end
+
+      it "rolls out relatively evenly to 3 behaviors" do
+        subject_experiment.variant(:variant3) {} # rubocop:disable Lint/EmptyBlock
+
+        100.times { |i| run_cycle(subject_experiment, value: i) }
+
+        expect(counts).to eq(variant1: 31, variant2: 29, variant3: 40)
+      end
+
+      context "when distribution is specified as an array" do
+        before do
+          subject_experiment.rollout(described_class, distribution: [32, 25, 43])
+        end
+
+        it "rolls out with the expected distribution" do
+          subject_experiment.variant(:variant3) {} # rubocop:disable Lint/EmptyBlock
+
+          100.times { |i| run_cycle(subject_experiment, value: i) }
+
+          expect(counts).to eq(variant1: 39, variant2: 24, variant3: 37)
+        end
+      end
+
+      context "when distribution is specified as a hash" do
+        before do
+          subject_experiment.rollout(described_class, distribution: { variant1: 90, variant2: 10 })
+        end
+
+        it "rolls out with the expected distribution" do
+          100.times { |i| run_cycle(subject_experiment, value: i) }
+
+          expect(counts).to eq(variant1: 95, variant2: 5)
+        end
+      end
+
+      def run_cycle(experiment, **context)
+        experiment.instance_variable_set(:@_assigned_variant_name, nil)
+        experiment.context(context) if context
+
+        begin
+          experiment.cache.delete
+        rescue StandardError
+          nil
+        end
 
-      expect(subject.execute_assignment).to eq(:variant2)
+        counts[experiment.assigned.name] += 1
+      end
     end
   end