Add a test for the add_user for optional context

Add experiment record context for growth experiment users

app/models/experiment.rb

@@ -5,8 +5,8 @@ class Experiment < ApplicationRecord
 
   validates :name, presence: true, uniqueness: true, length: { maximum: 255 }
 
-  def self.add_user(name, group_type, user)
-    find_or_create_by!(name: name).record_user_and_group(user, group_type)
+  def self.add_user(name, group_type, user, context = {})
+    find_or_create_by!(name: name).record_user_and_group(user, group_type, context)
   end
 
   def self.record_conversion_event(name, user)
@@ -14,8 +14,8 @@ class Experiment < ApplicationRecord
   end
 
   # 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)
+  def record_user_and_group(user, group_type, context = {})
+    experiment_users.find_or_initialize_by(user: user).update!(group_type: group_type, context: context)
   end
 
   def record_conversion_event_for_user(user)

changelogs/unreleased/234013-add-experiment-record-context.yml

+---
+title: Add context to the experiment user records
+merge_request: 48896
+author:
+type: added

db/migrate/20201201190002_add_other_context_to_experiment_user.rb

+# frozen_string_literal: true
+
+class AddOtherContextToExperimentUser < ActiveRecord::Migration[6.0]
+  include Gitlab::Database::MigrationHelpers
+
+  DOWNTIME = false
+
+  def up
+    with_lock_retries do
+      add_column :experiment_users, :context, :jsonb, default: {}, null: false
+    end
+  end
+
+  def down
+    with_lock_retries do
+      remove_column :experiment_users, :context
+    end
+  end
+end

db/schema_migrations/20201201190002

+f4ec800e68cbe092775b428d3ff85a4a84be0d55d70e59d23de390847ea3c2b7
\ No newline at end of file

db/structure.sql

@@ -12124,7 +12124,8 @@ CREATE TABLE experiment_users (
     group_type smallint DEFAULT 0 NOT NULL,
     created_at timestamp with time zone NOT NULL,
     updated_at timestamp with time zone NOT NULL,
-    converted_at timestamp with time zone
+    converted_at timestamp with time zone,
+    context jsonb DEFAULT '{}'::jsonb NOT NULL
 );
 
 CREATE SEQUENCE experiment_users_id_seq

ee/spec/controllers/registrations/welcome_controller_spec.rb

@@ -144,7 +144,7 @@ RSpec.describe Registrations::WelcomeController do
 
             with_them do
               it 'adds the user to the experiments table with the correct group_type' do
-                expect(::Experiment).to receive(:add_user).with(:onboarding_issues, group_type, user)
+                expect(::Experiment).to receive(:add_user).with(:onboarding_issues, group_type, user, {})
 
                 subject
               end

lib/gitlab/experimentation/controller_concern.rb

@@ -61,11 +61,11 @@ module Gitlab
         end
       end
 
-      def record_experiment_user(experiment_key)
+      def record_experiment_user(experiment_key, context = {})
         return if dnt_enabled?
         return unless Experimentation.active?(experiment_key) && current_user
 
-        ::Experiment.add_user(experiment_key, tracking_group(experiment_key, nil, subject: current_user), current_user)
+        ::Experiment.add_user(experiment_key, tracking_group(experiment_key, nil, subject: current_user), current_user, context)
       end
 
       def record_experiment_conversion_event(experiment_key)

spec/db/schema_spec.rb

@@ -184,6 +184,7 @@ RSpec.describe 'Database schema' do
     "ApplicationSetting" => %w[repository_storages_weighted],
     "AlertManagement::Alert" => %w[payload],
     "Ci::BuildMetadata" => %w[config_options config_variables],
+    "ExperimentUser" => %w[context],
     "Geo::Event" => %w[payload],
     "GeoNodeStatus" => %w[status],
     "Operations::FeatureFlagScope" => %w[strategies],

spec/lib/gitlab/experimentation/controller_concern_spec.rb

@@ -392,6 +392,7 @@ RSpec.describe Gitlab::Experimentation::ControllerConcern, type: :controller do
 
   describe '#record_experiment_user' do
     let(:user) { build(:user) }
+    let(:context) { { a: 42 } }
 
     context 'when the experiment is enabled' do
       before do
@@ -405,9 +406,9 @@ RSpec.describe Gitlab::Experimentation::ControllerConcern, type: :controller do
         end
 
         it 'calls add_user on the Experiment model' do
-          expect(::Experiment).to receive(:add_user).with(:test_experiment, :experimental, user)
+          expect(::Experiment).to receive(:add_user).with(:test_experiment, :experimental, user, context)
 
-          controller.record_experiment_user(:test_experiment)
+          controller.record_experiment_user(:test_experiment, context)
         end
       end
 
@@ -417,9 +418,9 @@ RSpec.describe Gitlab::Experimentation::ControllerConcern, type: :controller do
         end
 
         it 'calls add_user on the Experiment model' do
-          expect(::Experiment).to receive(:add_user).with(:test_experiment, :control, user)
+          expect(::Experiment).to receive(:add_user).with(:test_experiment, :control, user, context)
 
-          controller.record_experiment_user(:test_experiment)
+          controller.record_experiment_user(:test_experiment, context)
         end
       end
     end
@@ -433,7 +434,7 @@ RSpec.describe Gitlab::Experimentation::ControllerConcern, type: :controller do
       it 'does not call add_user on the Experiment model' do
         expect(::Experiment).not_to receive(:add_user)
 
-        controller.record_experiment_user(:test_experiment)
+        controller.record_experiment_user(:test_experiment, context)
       end
     end
 
@@ -445,7 +446,7 @@ RSpec.describe Gitlab::Experimentation::ControllerConcern, type: :controller do
       it 'does not call add_user on the Experiment model' do
         expect(::Experiment).not_to receive(:add_user)
 
-        controller.record_experiment_user(:test_experiment)
+        controller.record_experiment_user(:test_experiment, context)
       end
     end
 
@@ -462,9 +463,9 @@ RSpec.describe Gitlab::Experimentation::ControllerConcern, type: :controller do
         end
 
         it 'calls add_user on the Experiment model' do
-          expect(::Experiment).to receive(:add_user).with(:test_experiment, :control, user)
+          expect(::Experiment).to receive(:add_user).with(:test_experiment, :control, user, context)
 
-          controller.record_experiment_user(:test_experiment)
+          controller.record_experiment_user(:test_experiment, context)
         end
       end
 
@@ -476,7 +477,7 @@ RSpec.describe Gitlab::Experimentation::ControllerConcern, type: :controller do
         it 'does not call add_user on the Experiment model' do
           expect(::Experiment).not_to receive(:add_user)
 
-          controller.record_experiment_user(:test_experiment)
+          controller.record_experiment_user(:test_experiment, context)
         end
       end
     end

spec/models/experiment_spec.rb

@@ -19,20 +19,21 @@ RSpec.describe Experiment do
     let_it_be(:experiment_name) { :experiment_key }
     let_it_be(:user) { 'a user' }
     let_it_be(:group) { 'a group' }
+    let_it_be(:context) { { a: 42 } }
 
-    subject(:add_user) { described_class.add_user(experiment_name, group, user) }
+    subject(:add_user) { described_class.add_user(experiment_name, group, user, context) }
 
     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)
+          allow(experiment).to receive(:record_user_and_group).with(user, group, context)
         end
         expect { add_user }.to change(described_class, :count).by(1)
       end
 
-      it 'forwards the user and group_type to the instance' do
+      it 'forwards the user, group_type, and context to the instance' do
         expect_next_instance_of(described_class) do |experiment|
-          expect(experiment).to receive(:record_user_and_group).with(user, group)
+          expect(experiment).to receive(:record_user_and_group).with(user, group, context)
         end
         add_user
       end
@@ -43,18 +44,26 @@ RSpec.describe Experiment do
 
       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)
+          allow(experiment).to receive(:record_user_and_group).with(user, group, context)
         end
         expect { add_user }.not_to change(described_class, :count)
       end
 
-      it 'forwards the user and group_type to the instance' do
+      it 'forwards the user, group_type, and context to the instance' do
         expect_next_found_instance_of(described_class) do |experiment|
-          expect(experiment).to receive(:record_user_and_group).with(user, group)
+          expect(experiment).to receive(:record_user_and_group).with(user, group, context)
         end
         add_user
       end
     end
+
+    it 'works without the optional context argument' do
+      allow_next_instance_of(described_class) do |experiment|
+        expect(experiment).to receive(:record_user_and_group).with(user, group, {})
+      end
+
+      expect { described_class.add_user(experiment_name, group, user) }.not_to raise_error
+    end
   end
 
   describe '.record_conversion_event' do
@@ -131,8 +140,9 @@ RSpec.describe Experiment do
     let_it_be(:user) { create(:user) }
 
     let(:group) { :control }
+    let(:context) { { a: 42 } }
 
-    subject(:record_user_and_group) { experiment.record_user_and_group(user, group) }
+    subject(:record_user_and_group) { experiment.record_user_and_group(user, group, context) }
 
     context 'when an experiment_user does not yet exist for the given user' do
       it 'creates a new experiment_user record' do
@@ -143,24 +153,35 @@ RSpec.describe Experiment do
         record_user_and_group
         expect(ExperimentUser.last.group_type).to eq('control')
       end
+
+      it 'adds the correct context to the experiment_user' do
+        record_user_and_group
+        expect(ExperimentUser.last.context).to eq({ 'a' => 42 })
+      end
     end
 
     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)
+        experiment.record_user_and_group(user, :control, context)
       end
 
       it 'does not create a new experiment_user record' do
         expect { record_user_and_group }.not_to change(ExperimentUser, :count)
       end
 
-      context 'but the group_type has changed' do
+      context 'but the group_type and context has changed' do
         let(:group) { :experimental }
+        let(:context) { { b: 37 } }
 
-        it 'updates the existing experiment_user record' do
+        it 'updates the existing experiment_user record with group_type' do
           expect { record_user_and_group }.to change { ExperimentUser.last.group_type }
         end
+
+        it 'updates the existing experiment_user record with context' do
+          record_user_and_group
+          expect(ExperimentUser.last.context).to eq({ 'b' => 37 })
+        end
       end
     end
   end