Add add_group method to Experiment model

app/models/experiment.rb

@@ -10,6 +10,10 @@ class Experiment < ApplicationRecord
     find_or_create_by!(name: name).record_user_and_group(user, group_type, context)
   end
 
+  def self.add_group(name, variant:, group:)
+    find_or_create_by!(name: name).record_group_and_variant!(group, variant)
+  end
+
   def self.record_conversion_event(name, user)
     find_or_create_by!(name: name).record_conversion_event_for_user(user)
   end
@@ -24,4 +28,8 @@ class Experiment < ApplicationRecord
   def record_conversion_event_for_user(user)
     experiment_users.find_by(user: user, converted_at: nil)&.touch(:converted_at)
   end
+
+  def record_group_and_variant!(group, variant)
+    experiment_subjects.find_or_initialize_by(group: group).update!(variant: variant)
+  end
 end

spec/models/experiment_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe Experiment do
+  include AfterNextHelpers
+
   subject { build(:experiment) }
 
   describe 'associations' do
@@ -67,6 +69,33 @@ RSpec.describe Experiment do
     end
   end
 
+  describe '.add_group' do
+    let_it_be(:experiment_name) { :experiment_key }
+    let_it_be(:variant) { :control }
+    let_it_be(:group) { build(:group) }
+
+    subject(:add_group) { described_class.add_group(experiment_name, variant: variant, group: group) }
+
+    context 'when an experiment with the provided name does not exist' do
+      it 'creates a new experiment record' do
+        allow_next(described_class, name: :experiment_key)
+          .to receive(:record_group_and_variant!).with(group, variant)
+
+        expect { add_group }.to change(described_class, :count).by(1)
+      end
+    end
+
+    context 'when an experiment with the provided name already exists' do
+      before do
+        create(:experiment, name: experiment_name)
+      end
+
+      it 'does not create a new experiment record' do
+        expect { add_group }.not_to change(described_class, :count)
+      end
+    end
+  end
+
   describe '.record_conversion_event' do
     let_it_be(:user) { build(:user) }
 
@@ -136,6 +165,34 @@ RSpec.describe Experiment do
     end
   end
 
+  describe '#record_group_and_variant!' do
+    let_it_be(:group) { create(:group) }
+    let_it_be(:variant) { :control }
+    let_it_be(:experiment) { create(:experiment) }
+
+    subject(:record_group_and_variant!) { experiment.record_group_and_variant!(group, variant) }
+
+    context 'when no existing experiment_subject record exists for the given group' do
+      it 'creates an experiment_subject record' do
+        expect_next(ExperimentSubject).to receive(:update!).with(variant: variant).and_call_original
+
+        expect { record_group_and_variant! }.to change(ExperimentSubject, :count).by(1)
+      end
+    end
+
+    context 'when an existing experiment_subject exists for the given group' do
+      context 'but it belonged to a different variant' do
+        let!(:experiment_subject) do
+          create(:experiment_subject, experiment: experiment, group: group, user: nil, variant: :experimental)
+        end
+
+        it 'updates the variant value' do
+          expect { record_group_and_variant! }.to change { experiment_subject.reload.variant }.to('control')
+        end
+      end
+    end
+  end
+
   describe '#record_user_and_group' do
     let_it_be(:experiment) { create(:experiment) }
     let_it_be(:user) { create(:user) }