lib/field_test/experiment.rb in field_test-0.1.1 vs lib/field_test/experiment.rb in field_test-0.1.2

@@ -1,30 +1,32 @@
 module FieldTest
   class Experiment
-    attr_reader :id, :name, :variants, :winner, :started_at, :ended_at
+    attr_reader :id, :name, :variants, :weights, :winner, :started_at, :ended_at
 
     def initialize(attributes)
       attributes = attributes.symbolize_keys
       @id = attributes[:id]
       @name = attributes[:name] || @id.to_s.titleize
       @variants = attributes[:variants]
+      @weights = @variants.size.times.map { |i| attributes[:weights].to_a[i] || 1 }
       @winner = attributes[:winner]
       @started_at = Time.zone.parse(attributes[:started_at].to_s) if attributes[:started_at]
       @ended_at = Time.zone.parse(attributes[:ended_at].to_s) if attributes[:ended_at]
     end
 
     def variant(participants, options = {})
       return winner if winner
       return variants.first if options[:exclude]
 
       participants = FieldTest::Participant.standardize(participants)
+      check_participants(participants)
       membership = membership_for(participants) || FieldTest::Membership.new(experiment: id)
 
       if options[:variant] && variants.include?(options[:variant])
         membership.variant = options[:variant]
       else
-        membership.variant ||= variants.sample
+        membership.variant ||= weighted_variant
       end
 
       # upgrade to preferred participant
       membership.participant = participants.first
 
@@ -39,10 +41,11 @@
       membership.try(:variant) || variants.first
     end
 
     def convert(participants)
       participants = FieldTest::Participant.standardize(participants)
+      check_participants(participants)
       membership = membership_for(participants)
 
       if membership
         membership.converted = true
         membership.save! if membership.changed?
@@ -115,12 +118,27 @@
       end
     end
 
     private
 
+      def check_participants(participants)
+        raise FieldTest::UnknownParticipant, "Use the :participant option to specify a participant" if participants.empty?
+      end
+
       def membership_for(participants)
         memberships = self.memberships.where(participant: participants).index_by(&:participant)
         participants.map { |part| memberships[part] }.compact.first
+      end
+
+      def weighted_variant
+        total = weights.sum.to_f
+        pick = rand
+        n = 0
+        weights.map { |w| w / total }.each_with_index do |w, i|
+          n += w
+          return variants[i] if n >= pick
+        end
+        variants.last
       end
 
       # formula from
       # http://www.evanmiller.org/bayesian-ab-testing.html
       def prob_b_beats_a(alpha_a, beta_a, alpha_b, beta_b)