Invoke on_assignment callback after assignment is stored in database; add global on_assignment callback

lib/vanity/configuration.rb

@@ -49,6 +49,8 @@ def default_request_filter(request) # :nodoc:
       on_datastore_error: ->(error, klass, method, arguments) {
         default_on_datastore_error(error, klass, method, arguments)
       },
+      on_assignment: nil,
+      after_assignment: nil,
       request_filter: ->(request) { default_request_filter(request) },
       templates_path: File.expand_path(File.join(File.dirname(__FILE__), 'templates')),
       use_js: false,
@@ -181,6 +183,12 @@ def default_request_filter(request) # :nodoc:
     # Cookie path. If true, cookie will not be available to JS. By default false.
     attr_writer :cookie_httponly
 
+    # Default callback on assigment
+    attr_writer :on_assignment
+
+    # Default callback after assigment
+    attr_writer :after_assignment
+
     # We independently list each attr_accessor to includes docs, otherwise
     # something like DEFAULTS.each { |key, value| attr_accessor key } would be
     # shorter.

lib/vanity/experiment/ab_test.rb

@@ -616,13 +616,17 @@ def alternative_for(identity)
       # Saves the assignment of an alternative to a person and performs the
       # necessary housekeeping. Ignores repeat identities and filters using
       # Playground#request_filter.
-      def save_assignment(identity, index, request)
+      def save_assignment(identity, index, _request)
         return if index == connection.ab_showing(@id, identity)
+        return if connection.ab_seen @id, identity, index
 
         call_on_assignment_if_available(identity, index)
         rebalance_if_necessary!
 
         connection.ab_add_participant(@id, index, identity)
+
+        call_after_assignment_if_available(identity, index)
+
         check_completion!
       end
 
@@ -635,10 +639,29 @@ def call_on_assignment_if_available(identity, index)
         # if we have an on_assignment block, call it on new assignments
         if defined?(@on_assignment_block) && @on_assignment_block
           assignment = alternatives[index]
-          if !connection.ab_seen @id, identity, index
-            @on_assignment_block.call(Vanity.context, identity, assignment, self)
-          end
+          @on_assignment_block.call(Vanity.context, identity, assignment, self)
+
+          return
         end
+
+        return unless Vanity.configuration.on_assignment.is_a?(Proc)
+
+        Vanity.configuration.on_assignment.call(Vanity.context, identity, assignment, self)
+      end
+
+      def call_after_assignment_if_available(identity, index)
+        # if we have an after_assignment_block block, call it after new assignments
+        if defined?(@after_assignment_block) && @after_assignment_block
+          assignment = alternatives[index]
+
+          @after_assignment_block.call(Vanity.context, identity, assignment, self)
+
+          return
+        end
+
+        return unless Vanity.configuration.after_assignment.is_a?(Proc)
+
+        Vanity.configuration.after_assignment.call(Vanity.context, identity, assignment, self)
       end
 
       def rebalance_if_necessary!
@@ -687,10 +710,8 @@ def build_alternatives(args)
         # We're really only interested in 90%, 95%, 99% and 99.9%.
         Z_TO_PROBABILITY = [90, 95, 99, 99.9].map { |pct| [norm_dist.find { |x,a| a >= pct }.first, pct] }.reverse
       end
-
     end
 
-
     module Definition
       # Define an A/B test with the given name. For example:
       #   ab_test "New Banner" do
@@ -700,6 +721,5 @@ def ab_test(name, &block)
         define name, :ab_test, &block
       end
     end
-
   end
 end

lib/vanity/experiment/base.rb

@@ -41,6 +41,7 @@ def initialize(playground, id, name, options = nil)
         @options = options || {}
         @identify_block = method(:default_identify)
         @on_assignment_block = nil
+        @after_assignment_block = nil
       end
 
       # Human readable experiment name (first argument you pass when creating a
@@ -96,6 +97,20 @@ def on_assignment(&block)
         @on_assignment_block = block
       end
 
+      # Defines any additional actions to take when a new assignment was made for the current experiment
+      #
+      # For example
+      #   ab_test "Project widget" do
+      #     alternatives :small, :medium, :large
+      #     after_assignment do |controller, identity, assignment|
+      #       # Do something useful
+      #     end
+      #   end
+      def after_assignment(&block)
+        fail "Missing block" unless block
+        @after_assignment_block = block
+      end
+
       # -- Reporting --
 
       # Sets or returns description. For example

test/experiment/ab_test.rb

@@ -520,6 +520,44 @@ def test_returns_the_same_alternative_consistently_when_on_assignment_is_set
     end
   end
 
+  def test_calls_default_on_assignment
+    on_assignment_called_times = 0
+
+    Vanity.configuration.on_assignment = proc do |_controller, _identity, _assignment|
+      on_assignment_called_times += 1
+    end
+
+    new_ab_test :foobar do
+      alternatives "foo", "bar"
+      default "foo"
+      identify { "6e98ec" }
+      metrics :coolness
+    end
+
+    2.times { experiment(:foobar).chooses("foo") }
+    assert_equal 1, on_assignment_called_times
+  end
+
+  # -- after_assignment --
+
+  def test_calls_default_after_assignment
+    after_assignment_called_times = 0
+
+    Vanity.configuration.after_assignment = proc do |_controller, _identity, _assignment|
+      after_assignment_called_times += 1
+    end
+
+    new_ab_test :foobar do
+      alternatives "foo", "bar"
+      default "foo"
+      identify { "6e98ec" }
+      metrics :coolness
+    end
+
+    2.times { experiment(:foobar).chooses("foo") }
+    assert_equal 1, after_assignment_called_times
+  end
+
   # -- ab_assigned --
 
   def test_ab_assigned

test/playground_test.rb

@@ -130,5 +130,4 @@
       assert_equal [[Vanity.playground.experiment(:foobar), alt]], Vanity.playground.participant_info("abcdef")
     end
   end
-
 end