Update existing records when recording exp. users

Previously we just exited out of the user recording process if we found
an existing experiment_users record. Now we update that existing record.
This is very useful in the case that the rollout percentage changes and
a previously recorded "control" group user moves into the "experimental"
group.

Also, I removed a lot of the extra association setup from the Experiment
model as it wasn't being used anywhere outside of the model and my
updates to that model no longer make use of those associations.

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