Merge branch...

Merge branch '236185-update-experiment-grouping-on-subsequent-record_experiment_user-calls' into 'master'

Update experiment grouping on subsequent record_experiment_user calls

See merge request gitlab-org/gitlab!46390

app/models/experiment.rb

 # frozen_string_literal: true
 
 class Experiment < ApplicationRecord
-  include ::Gitlab::Experimentation::GroupTypes
-
   has_many :experiment_users
-  has_many :users, through: :experiment_users
-  has_many :control_group_users, -> { merge(ExperimentUser.control) }, through: :experiment_users, source: :user
-  has_many :experimental_group_users, -> { merge(ExperimentUser.experimental) }, through: :experiment_users, source: :user
 
   validates :name, presence: true, uniqueness: true, length: { maximum: 255 }
 
   def self.add_user(name, group_type, user)
-    experiment = find_or_create_by(name: name)
-
-    return unless experiment
-    return if experiment.experiment_users.where(user: user).exists?
-
-    group_type == GROUP_CONTROL ? experiment.add_control_user(user) : experiment.add_experimental_user(user)
-  end
+    return unless experiment = find_or_create_by(name: name)
 
-  def add_control_user(user)
-    control_group_users << user
+    experiment.record_user_and_group(user, group_type)
   end
 
-  def add_experimental_user(user)
-    experimental_group_users << user
+  # Create or update the recorded experiment_user row for the user in this experiment.
+  def record_user_and_group(user, group_type)
+    experiment_users.find_or_initialize_by(user: user).update!(group_type: group_type)
   end
 end

app/models/experiment_user.rb

@@ -8,5 +8,7 @@ class ExperimentUser < ApplicationRecord
 
   enum group_type: { GROUP_CONTROL => 0, GROUP_EXPERIMENTAL => 1 }
 
+  validates :experiment_id, presence: true
+  validates :user_id, presence: true
   validates :group_type, presence: true
 end

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