1 files changed, 53 insertions, 67 deletions

diff --git a/spec/models/experiment_spec.rb b/spec/models/experiment_spec.rb
index 64cd2da4621..587f410c9be 100644
--- a/spec/models/experiment_spec.rb
+++ b/spec/models/experiment_spec.rb
@@ -7,32 +7,6 @@ RSpec.describe Experiment do
 
   describe 'associations' do
     it { is_expected.to have_many(:experiment_users) }
-    it { is_expected.to have_many(:users) }
-    it { is_expected.to have_many(:control_group_users) }
-    it { is_expected.to have_many(:experimental_group_users) }
-
-    describe 'control_group_users and experimental_group_users' do
-      let(:experiment) { create(:experiment) }
-      let(:control_group_user) { build(:user) }
-      let(:experimental_group_user) { build(:user) }
-
-      before do
-        experiment.control_group_users << control_group_user
-        experiment.experimental_group_users << experimental_group_user
-      end
-
-      describe 'control_group_users' do
-        subject { experiment.control_group_users }
-
-        it { is_expected.to contain_exactly(control_group_user) }
-      end
-
-      describe 'experimental_group_users' do
-        subject { experiment.experimental_group_users }
-
-        it { is_expected.to contain_exactly(experimental_group_user) }
-      end
-    end
   end
 
   describe 'validations' do
@@ -42,71 +16,83 @@ RSpec.describe Experiment do
   end
 
   describe '.add_user' do
-    let(:name) { :experiment_key }
-    let(:user) { build(:user) }
+    let_it_be(:experiment_name) { :experiment_key }
+    let_it_be(:user) { 'a user' }
+    let_it_be(:group) { 'a group' }
 
-    let!(:experiment) { create(:experiment, name: name) }
+    subject(:add_user) { described_class.add_user(experiment_name, group, user) }
 
-    subject { described_class.add_user(name, :control, user) }
-
-    describe 'creating a new experiment record' do
-      context 'an experiment with the provided name already exists' do
-        it 'does not create a new experiment record' do
-          expect { subject }.not_to change(Experiment, :count)
+    context 'when an experiment with the provided name does not exist' do
+      it 'creates a new experiment record' do
+        allow_next_instance_of(described_class) do |experiment|
+          allow(experiment).to receive(:record_user_and_group).with(user, group)
         end
+        expect { add_user }.to change(described_class, :count).by(1)
       end
 
-      context 'an experiment with the provided name does not exist yet' do
-        let(:experiment) { nil }
-
-        it 'creates a new experiment record' do
-          expect { subject }.to change(Experiment, :count).by(1)
+      it 'forwards the user and group_type to the instance' do
+        expect_next_instance_of(described_class) do |experiment|
+          expect(experiment).to receive(:record_user_and_group).with(user, group)
         end
+        add_user
       end
     end
 
-    describe 'creating a new experiment_user record' do
-      context 'an experiment_user record for this experiment already exists' do
-        before do
-          subject
-        end
+    context 'when an experiment with the provided name already exists' do
+      let_it_be(:experiment) { create(:experiment, name: experiment_name) }
 
-        it 'does not create a new experiment_user record' do
-          expect { subject }.not_to change(ExperimentUser, :count)
+      it 'does not create a new experiment record' do
+        allow_next_found_instance_of(described_class) do |experiment|
+          allow(experiment).to receive(:record_user_and_group).with(user, group)
         end
+        expect { add_user }.not_to change(described_class, :count)
       end
 
-      context 'an experiment_user record for this experiment does not exist yet' do
-        it 'creates a new experiment_user record' do
-          expect { subject }.to change(ExperimentUser, :count).by(1)
-        end
-
-        it 'assigns the correct group_type to the experiment_user' do
-          expect { subject }.to change { experiment.control_group_users.count }.by(1)
+      it 'forwards the user and group_type to the instance' do
+        expect_next_found_instance_of(described_class) do |experiment|
+          expect(experiment).to receive(:record_user_and_group).with(user, group)
         end
+        add_user
       end
     end
   end
 
-  describe '#add_control_user' do
-    let(:experiment) { create(:experiment) }
-    let(:user) { build(:user) }
+  describe '#record_user_and_group' do
+    let_it_be(:experiment) { create(:experiment) }
+    let_it_be(:user) { create(:user) }
 
-    subject { experiment.add_control_user(user) }
+    let(:group) { :control }
 
-    it 'creates a new experiment_user record and assigns the correct group_type' do
-      expect { subject }.to change { experiment.control_group_users.count }.by(1)
+    subject(:record_user_and_group) { experiment.record_user_and_group(user, group) }
+
+    context 'when an experiment_user does not yet exist for the given user' do
+      it 'creates a new experiment_user record' do
+        expect { record_user_and_group }.to change(ExperimentUser, :count).by(1)
+      end
+
+      it 'assigns the correct group_type to the experiment_user' do
+        record_user_and_group
+        expect(ExperimentUser.last.group_type).to eq('control')
+      end
     end
-  end
 
-  describe '#add_experimental_user' do
-    let(:experiment) { create(:experiment) }
-    let(:user) { build(:user) }
+    context 'when an experiment_user already exists for the given user' do
+      before do
+        # Create an existing experiment_user for this experiment and the :control group
+        experiment.record_user_and_group(user, :control)
+      end
+
+      it 'does not create a new experiment_user record' do
+        expect { record_user_and_group }.not_to change(ExperimentUser, :count)
+      end
 
-    subject { experiment.add_experimental_user(user) }
+      context 'but the group_type has changed' do
+        let(:group) { :experimental }
 
-    it 'creates a new experiment_user record and assigns the correct group_type' do
-      expect { subject }.to change { experiment.experimental_group_users.count }.by(1)
+        it 'updates the existing experiment_user record' do
+          expect { record_user_and_group }.to change { ExperimentUser.last.group_type }
+        end
+      end
     end
   end
 end