1 files changed, 5 insertions, 14 deletions
diff --git a/app/models/experiment.rb b/app/models/experiment.rb
index 25640385536..f179a1fc6ce 100644
--- a/app/models/experiment.rb
+++ b/app/models/experiment.rb
@@ -2,26 +2,17 @@
 
 class Experiment < ApplicationRecord
   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 = find_or_create_by(name: name)
 
-    return unless experiment
-    return if experiment.experiment_users.where(user: user).exists?
-
-    group_type == ::Gitlab::Experimentation::GROUP_CONTROL ? experiment.add_control_user(user) : experiment.add_experimental_user(user)
-  end
-
-  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