gitlab-experiment 0.6.5 → 0.7.0

data/lib/gitlab/experiment/rspec.rb

@@ -6,226 +6,293 @@ module Gitlab
       autoload :Trackable, 'gitlab/experiment/test_behaviors/trackable.rb'
     end
 
-    module RSpecHelpers
-      def stub_experiments(experiments, times = nil)
-        experiments.each { |experiment| wrapped_experiment(experiment, times) }
+    WrappedExperiment = Struct.new(:klass, :experiment_name, :variant_name, :expectation_chain, :blocks)
+
+    module RSpecMocks
+      @__gitlab_experiment_receivers = {}
+
+      def self.track_gitlab_experiment_receiver(method, receiver)
+        # Leverage the `>=` method on Gitlab::Experiment to determine if the receiver is an experiment, not the other
+        # way round -- `receiver.<=` could be mocked and we want to be extra careful.
+        (@__gitlab_experiment_receivers[method.to_s] ||= []) << receiver if Gitlab::Experiment >= receiver
+      rescue StandardError # again, let's just be extra careful
+        false
       end
 
-      def wrapped_experiment(experiment, times = nil, expected = false, &block)
-        klass, experiment_name, variant_name = *experiment_details(experiment)
-        base_klass = Configuration.base_class.constantize
+      def self.bind_gitlab_experiment_receiver(method)
+        method.unbind.bind(@__gitlab_experiment_receivers[method.to_s].pop)
+      end
 
-        # Set expectations on experiment classes so we can and_wrap_original with more specific args
-        experiment_klasses = base_klass.descendants.reject { |k| k == klass }
-        experiment_klasses.push(base_klass).each do |k|
-          allow(k).to receive(:new).and_call_original
+      module MethodDouble
+        def proxy_method_invoked(receiver, *args, &block)
+          RSpecMocks.track_gitlab_experiment_receiver(original_method, receiver)
+          super
         end
+      end
+    end
 
-        receiver = receive(:new)
-
-        # Be specific for BaseClass calls
-        receiver = receiver.with(experiment_name, any_args) if experiment_name && klass == base_klass
+    module RSpecHelpers
+      def stub_experiments(experiments)
+        experiments.each do |experiment|
+          wrapped_experiment(experiment, remock: true) do |instance, wrapped|
+            # Stub internal methods that will make it behave as we've instructed.
+            allow(instance).to receive(:enabled?) { wrapped.variant_name != false }
+            if instance.respond_to?(:experiment_group?, true)
+              allow(instance).to receive(:experiment_group?) { !(wrapped.variant_name == false) }
+            end
 
-        receiver.exactly(times).times if times
+            # Stub the variant resolution logic to handle true/false, and named variants.
+            allow(instance).to receive(:resolve_variant_name).and_wrap_original { |method|
+              # Call the original method if we specified simply `true`.
+              wrapped.variant_name == true ? method.call : wrapped.variant_name
+            }
+          end
+        end
 
-        # Set expectations on experiment class of interest
-        allow_or_expect_klass = expected ? expect(klass) : allow(klass)
-        allow_or_expect_klass.to receiver.and_wrap_original do |method, *original_args, &original_block|
-          method.call(*original_args).tap do |e|
-            # Stub internal methods before calling the original_block
-            allow(e).to receive(:enabled?).and_return(true)
+        wrapped_experiments
+      end
 
-            if variant_name == true # passing true allows the rollout to do its job
-              allow(e).to receive(:experiment_group?).and_return(true)
-            else
-              allow(e).to receive(:resolve_variant_name).and_return(variant_name.to_s)
-            end
+      def wrapped_experiment(experiment, remock: false, &block)
+        klass, experiment_name, variant_name = *extract_experiment_details(experiment)
 
-            # Stub/set expectations before calling the original_block
-            yield e if block
+        wrapped_experiment = wrapped_experiments[experiment_name] =
+          (!remock && wrapped_experiments[experiment_name]) ||
+          WrappedExperiment.new(klass, experiment_name, variant_name, wrapped_experiment_chain_for(klass), [])
 
-            original_block.call(e) if original_block.present?
-          end
-        end
+        wrapped_experiment.blocks << block if block
+        wrapped_experiment
       end
 
       private
 
-      def experiment_details(experiment)
-        if experiment.is_a?(Symbol)
-          experiment_name = experiment
-          variant_name = nil
+      def wrapped_experiments
+        @__wrapped_experiments ||= defined?(HashWithIndifferentAccess) ? HashWithIndifferentAccess.new : {}
+      end
+
+      def wrapped_experiment_chain_for(klass)
+        @__wrapped_experiment_chains ||= {}
+        @__wrapped_experiment_chains[klass.name || klass.object_id] ||= begin
+          allow(klass).to receive(:new).and_wrap_original do |method, *args, &original_block|
+            RSpecMocks.bind_gitlab_experiment_receiver(method).call(*args).tap do |instance|
+              wrapped = @__wrapped_experiments[instance.instance_variable_get(:@_name)]
+              wrapped&.blocks&.each { |b| b.call(instance, wrapped) }
+
+              original_block&.call(instance)
+            end
+          end
         end
+      end
+
+      def extract_experiment_details(experiment)
+        experiment_name = nil
+        variant_name = nil
 
+        experiment_name = experiment if experiment.is_a?(Symbol)
         experiment_name, variant_name = *experiment if experiment.is_a?(Array)
 
         base_klass = Configuration.base_class.constantize
-        variant_name = experiment.variant.name if experiment.is_a?(base_klass)
+        variant_name = experiment.assigned.name if experiment.is_a?(base_klass)
 
-        if experiment.class.name.nil? # Anonymous class instance
-          klass = experiment.class
-        elsif experiment.instance_of?(Class) # Class level stubbing, eg. "MyExperiment"
-          klass = experiment
-        else
-          experiment_name ||= experiment.instance_variable_get(:@name)
-          klass = base_klass.constantize(experiment_name)
-        end
+        resolved_klass = experiment_klass(experiment) { base_klass.constantize(experiment_name) }
+        experiment_name ||= experiment.instance_variable_get(:@_name)
 
-        if experiment_name && klass == base_klass
-          experiment_name = experiment_name.to_sym
+        [resolved_klass, experiment_name.to_s, variant_name]
+      end
 
-          # For experiment names like: "group/experiment-name"
-          experiment_name = experiment_name.to_s if experiment_name.inspect.include?('"')
+      def experiment_klass(experiment, &block)
+        if experiment.class.name.nil? # anonymous class instance
+          experiment.class
+        elsif experiment.instance_of?(Class) # class level stubbing, eg. "MyExperiment"
+          experiment
+        elsif block
+          yield
         end
-
-        [klass, experiment_name, variant_name]
       end
     end
 
     module RSpecMatchers
       extend RSpec::Matchers::DSL
 
-      def require_experiment(experiment, matcher_name, classes: false)
+      def require_experiment(experiment, matcher, instances_only: true)
         klass = experiment.instance_of?(Class) ? experiment : experiment.class
-        unless klass <= Gitlab::Experiment
-          raise(
-            ArgumentError,
-            "#{matcher_name} matcher is limited to experiment instances#{classes ? ' and classes' : ''}"
-          )
+        raise ArgumentError, "the #{matcher} matcher is limited to experiments" unless klass <= Gitlab::Experiment
+
+        if instances_only && experiment == klass
+          raise ArgumentError, "the #{matcher} matcher is limited to experiment instances"
+        end
+
+        experiment
+      end
+
+      matcher :register_behavior do |behavior_name|
+        match do |experiment|
+          @experiment = require_experiment(experiment, 'register_behavior')
+
+          block = @experiment.behaviors[behavior_name.to_s]
+          @return_expected = false unless block
+
+          if @return_expected
+            @actual_return = block.call
+            @expected_return == @actual_return
+          else
+            block
+          end
         end
 
-        if experiment == klass && !classes
-          raise ArgumentError, "#{matcher_name} matcher requires an instance of an experiment"
+        chain :with do |expected|
+          @return_expected = true
+          @expected_return = expected
         end
 
-        experiment != klass
+        failure_message do
+          add_details("expected the #{behavior_name} behavior to be registered")
+        end
+
+        failure_message_when_negated do
+          add_details("expected the #{behavior_name} behavior not to be registered")
+        end
+
+        def add_details(base)
+          details = []
+
+          if @return_expected
+            base = "#{base} with a return value"
+            details << "    expected return: #{@expected_return.inspect}\n" \
+                       "      actual return: #{@actual_return.inspect}"
+          else
+            details << "    behaviors: #{@experiment.behaviors.keys.inspect}"
+          end
+
+          details.unshift(base).join("\n")
+        end
       end
 
       matcher :exclude do |context|
-        ivar = :'@excluded'
-
         match do |experiment|
-          require_experiment(experiment, 'exclude')
-          experiment.context(context)
+          @experiment = require_experiment(experiment, 'exclude')
+          @experiment.context(context)
+          @experiment.instance_variable_set(:@_excluded, nil)
 
-          experiment.instance_variable_set(ivar, nil)
-          !experiment.run_callbacks(:exclusion_check) { :not_excluded }
+          !@experiment.run_callbacks(:exclusion_check) { :not_excluded }
         end
 
         failure_message do
-          %(expected #{context} to be excluded)
+          "expected #{context} to be excluded"
         end
 
         failure_message_when_negated do
-          %(expected #{context} not to be excluded)
+          "expected #{context} not to be excluded"
         end
       end
 
       matcher :segment do |context|
-        ivar = :'@variant_name'
-
         match do |experiment|
-          require_experiment(experiment, 'segment')
-          experiment.context(context)
+          @experiment = require_experiment(experiment, 'segment')
+          @experiment.context(context)
+          @experiment.instance_variable_set(:@_assigned_variant_name, nil)
+          @experiment.run_callbacks(:segmentation)
 
-          experiment.instance_variable_set(ivar, nil)
-          experiment.run_callbacks(:segmentation_check)
-
-          @actual = experiment.instance_variable_get(ivar)
-          @expected ? @actual.to_s == @expected.to_s : @actual.present?
+          @actual_variant = @experiment.instance_variable_get(:@_assigned_variant_name)
+          @expected_variant ? @actual_variant.to_s == @expected_variant.to_s : @actual_variant.present?
         end
 
         chain :into do |expected|
           raise ArgumentError, 'variant name must be provided' if expected.blank?
 
-          @expected = expected.to_s
+          @expected_variant = expected.to_s
         end
 
         failure_message do
-          %(expected #{context} to be segmented#{message_details})
+          add_details("expected #{context} to be segmented")
         end
 
         failure_message_when_negated do
-          %(expected #{context} not to be segmented#{message_details})
+          add_details("expected #{context} not to be segmented")
         end
 
-        def message_details
-          message = ''
-          message += %( into variant\n    expected variant: #{@expected}) if @expected
-          message += %(\n      actual variant: #{@actual}) if @actual
-          message
+        def add_details(base)
+          details = []
+
+          if @expected_variant
+            base = "#{base} into variant"
+            details << "    expected variant: #{@expected_variant.inspect}\n" \
+                       "      actual variant: #{@actual_variant.inspect}"
+          end
+
+          details.unshift(base).join("\n")
         end
       end
 
       matcher :track do |event, *event_args|
         match do |experiment|
-          expect_tracking_on(experiment, false, event, *event_args)
+          @experiment = require_experiment(experiment, 'track', instances_only: false)
+
+          set_expectations(event, *event_args, negated: false)
         end
 
         match_when_negated do |experiment|
-          expect_tracking_on(experiment, true, event, *event_args)
+          @experiment = require_experiment(experiment, 'track', instances_only: false)
+
+          set_expectations(event, *event_args, negated: true)
         end
 
-        chain :for do |expected_variant|
+        chain(:for) do |expected|
           raise ArgumentError, 'variant name must be provided' if expected.blank?
 
-          @expected_variant = expected_variant.to_s
+          @expected_variant = expected.to_s
         end
 
-        chain(:with_context) { |expected_context| @expected_context = expected_context }
+        chain(:with_context) do |expected|
+          raise ArgumentError, 'context name must be provided' if expected.nil?
 
-        chain(:on_next_instance) { @on_next_instance = true }
+          @expected_context = expected
+        end
 
-        def expect_tracking_on(experiment, negated, event, *event_args)
-          klass = experiment.instance_of?(Class) ? experiment : experiment.class
-          unless klass <= Gitlab::Experiment
-            raise(
-              ArgumentError,
-              "track matcher is limited to experiment instances and classes"
-            )
-          end
+        chain(:on_next_instance) { @on_next_instance = true }
 
+        def set_expectations(event, *event_args, negated:)
+          failure_message = failure_message_with_details(event, negated: negated)
           expectations = proc do |e|
-            @experiment = e
             allow(e).to receive(:track).and_call_original
 
             if negated
-              expect(e).not_to receive(:track).with(*[event, *event_args])
-            else
-              if @expected_variant
-                expect(@experiment.variant.name).to eq(@expected_variant), failure_message(:variant, event)
+              if @expected_variant || @expected_context
+                raise ArgumentError, 'cannot specify `for` or `with_context` when negating on tracking calls'
               end
 
-              if @expected_context
-                expect(@experiment.context.value).to include(@expected_context), failure_message(:context, event)
-              end
-
-              expect(e).to receive(:track).with(*[event, *event_args]).and_call_original
+              expect(e).not_to receive(:track).with(*[event, *event_args]), failure_message
+            else
+              expect(e.assigned.name).to(eq(@expected_variant), failure_message) if @expected_variant
+              expect(e.context.value).to(include(@expected_context), failure_message) if @expected_context
+              expect(e).to receive(:track).with(*[event, *event_args]).and_call_original, failure_message
             end
           end
 
-          if experiment.instance_of?(Class) || @on_next_instance
-            wrapped_experiment(experiment, nil, true) { |e| expectations.call(e) }
-          else
-            expectations.call(experiment)
-          end
+          return wrapped_experiment(@experiment, &expectations) if @on_next_instance || @experiment.instance_of?(Class)
+
+          expectations.call(@experiment)
         end
 
-        def failure_message(failure_type, event)
-          case failure_type
-          when :variant
-            <<~MESSAGE.strip
-              expected #{@experiment.inspect} to have tracked #{event.inspect} for variant
-                  expected variant: #{@expected_variant}
-                    actual variant: #{@experiment.variant.name}
-            MESSAGE
-          when :context
-            <<~MESSAGE.strip
-              expected #{@experiment.inspect} to have tracked #{event.inspect} with context
-                  expected context: #{@expected_context}
-                    actual context: #{@experiment.context.value}
-            MESSAGE
+        def failure_message_with_details(event, negated: false)
+          add_details("expected #{@experiment.inspect} #{negated ? 'not to' : 'to'} have tracked #{event.inspect}")
+        end
+
+        def add_details(base)
+          details = []
+
+          if @expected_variant
+            base = "#{base} for variant"
+            details << "    expected variant: #{@expected_variant.inspect}\n" \
+                       "      actual variant: #{@experiment.assigned.name.inspect})"
           end
+
+          if @expected_context
+            base = "#{base} with context"
+            details << "    expected context: #{@expected_context.inspect}\n" \
+                       "      actual context: #{@experiment.context.value.inspect})"
+          end
+
+          details.unshift(base).join("\n")
         end
       end
     end
@@ -248,4 +315,18 @@ RSpec.configure do |config|
   config.define_derived_metadata(file_path: Regexp.new('/spec/experiments/')) do |metadata|
     metadata[:type] = :experiment
   end
+
+  # We need to monkeypatch rspec-mocks because there's an issue around stubbing class methods that impacts us here.
+  #
+  # You can find out what the outcome is of the issues I've opened on rspec-mocks, and maybe some day this won't be
+  # needed.
+  #
+  # https://github.com/rspec/rspec-mocks/issues/1452
+  # https://github.com/rspec/rspec-mocks/issues/1451 (closed)
+  #
+  # The other way I've considered patching this is inside gitlab-experiment itself, by adding an Anonymous class and
+  # instantiating that instead of the configured base_class, and then it's less common but still possible to run into
+  # the issue.
+  require 'rspec/mocks/method_double'
+  RSpec::Mocks::MethodDouble.prepend(Gitlab::Experiment::RSpecMocks::MethodDouble)
 end

data/lib/gitlab/experiment/version.rb

@@ -2,6 +2,6 @@
 
 module Gitlab
   class Experiment
-    VERSION = '0.6.5'
+    VERSION = '0.7.0'
   end
 end