Adding callbacks

lib/factory_girl.rb

@@ -9,6 +9,7 @@
 require 'factory_girl/attribute/static'
 require 'factory_girl/attribute/dynamic'
 require 'factory_girl/attribute/association'
+require 'factory_girl/attribute/callback'
 require 'factory_girl/sequence'
 require 'factory_girl/aliases'
 

lib/factory_girl/attribute/callback.rb

@@ -0,0 +1,16 @@
+class Factory
+  class Attribute #:nodoc:
+
+    class Callback < Attribute  #:nodoc:
+      def initialize(name, block)
+        @name  = name.to_sym
+        @block = block
+      end
+
+      def add_to(proxy)
+        proxy.add_callback(name, @block)
+      end
+    end
+
+  end
+end

lib/factory_girl/factory.rb

@@ -4,6 +4,10 @@ class Factory
   class AssociationDefinitionError < RuntimeError
   end
 
+  # Raised when a callback is defined that has an invalid name
+  class InvalidCallbackNameError < RuntimeError
+  end
+
   class << self
     attr_accessor :factories #:nodoc:
 
@@ -184,6 +188,25 @@ def sequence (name, &block)
     add_attribute(name) { s.next }
   end
 
+  def after_build(&block)
+    callback(:after_build, &block)
+  end
+
+  def after_create(&block)
+    callback(:after_create, &block)
+  end
+
+  def after_stub(&block)
+    callback(:after_stub, &block)
+  end
+
+  def callback(name, &block)
+    unless [:after_build, :after_create, :after_stub].include?(name.to_sym)
+      raise InvalidCallbackNameError, "#{name} is not a valid callback name. Valid callback names are :after_build, :after_create, and :after_stub"
+    end
+    @attributes << Attribute::Callback.new(name.to_sym, block)
+  end
+
   # Generates and returns a Hash of attributes from this factory. Attributes
   # can be individually overridden by passing in a Hash of attribute => value
   # pairs.

lib/factory_girl/proxy.rb

@@ -1,6 +1,9 @@
 class Factory
 
   class Proxy #:nodoc:
+
+    attr_reader :callbacks
+
     def initialize(klass)
     end
 
@@ -14,6 +17,20 @@ def set(attribute, value)
     def associate(name, factory, attributes)
     end
 
+    def add_callback(name, block)
+      @callbacks ||= {}
+      @callbacks[name] ||= []
+      @callbacks[name] << block
+    end
+
+    def run_callbacks(name)
+      if @callbacks && @callbacks[name]
+        @callbacks[name].each do |block|
+          block.arity.zero? ? block.call : block.call(@instance)
+        end
+      end
+    end
+
     # Generates an association using the current build strategy.
     #
     # Arguments:

lib/factory_girl/proxy/build.rb

@@ -22,6 +22,7 @@ def association(factory, overrides = {})
       end
 
       def result
+        run_callbacks(:after_build)
         @instance
       end
     end

lib/factory_girl/proxy/create.rb

@@ -2,7 +2,9 @@ class Factory
   class Proxy #:nodoc:
     class Create < Build #:nodoc:
       def result
+        run_callbacks(:after_build)
         @instance.save!
+        run_callbacks(:after_create)
         @instance
       end
     end

lib/factory_girl/proxy/stub.rb

@@ -4,9 +4,9 @@ class Stub < Proxy #:nodoc:
       @@next_id = 1000
 
       def initialize(klass)
-        @stub = klass.new
-        @stub.id = next_id
-        @stub.instance_eval do
+        @instance = klass.new
+        @instance.id = next_id
+        @instance.instance_eval do
           def new_record?
             id.nil?
           end
@@ -26,11 +26,11 @@ def next_id
       end
 
       def get(attribute)
-        @stub.send(attribute)
+        @instance.send(attribute)
       end
 
       def set(attribute, value)
-        @stub.send(:"#{attribute}=", value)
+        @instance.send(:"#{attribute}=", value)
       end
 
       def associate(name, factory, attributes)
@@ -42,7 +42,8 @@ def association(factory, overrides = {})
       end
 
       def result
-        @stub
+        run_callbacks(:after_stub)
+        @instance
       end
     end
   end

spec/factory_girl/attribute/callback_spec.rb

@@ -0,0 +1,23 @@
+require File.expand_path(File.join(File.dirname(__FILE__), '..', '..', 'spec_helper'))
+
+describe Factory::Attribute::Callback do
+  before do
+    @name  = :after_create
+    @block = proc{ 'block' }
+    @attr  = Factory::Attribute::Callback.new(@name, @block)
+  end
+
+  it "should have a name" do
+    @attr.name.should == @name
+  end
+
+  it "should set its callback on a proxy" do
+    @proxy = "proxy"
+    mock(@proxy).add_callback(@name, @block)
+    @attr.add_to(@proxy)
+  end
+
+  it "should convert names to symbols" do
+    Factory::Attribute::Callback.new('name', nil).name.should == :name
+  end
+end

spec/factory_girl/factory_spec.rb

@@ -94,6 +94,38 @@
       end
     end
 
+    describe "adding a callback" do
+      it "should add a callback attribute when the after_build attribute is defined" do
+        mock(Factory::Attribute::Callback).new(:after_build, is_a(Proc)) { 'after_build callback' }
+        @factory.after_build {}
+        @factory.attributes.should include('after_build callback')
+      end
+
+      it "should add a callback attribute when the after_create attribute is defined" do
+        mock(Factory::Attribute::Callback).new(:after_create, is_a(Proc)) { 'after_create callback' }
+        @factory.after_create {}
+        @factory.attributes.should include('after_create callback')
+      end
+
+      it "should add a callback attribute when the after_stub attribute is defined" do
+        mock(Factory::Attribute::Callback).new(:after_stub, is_a(Proc)) { 'after_stub callback' }
+        @factory.after_stub {}
+        @factory.attributes.should include('after_stub callback')
+      end
+
+      it "should add a callback attribute when defining a callback" do
+        mock(Factory::Attribute::Callback).new(:after_create, is_a(Proc)) { 'after_create callback' }
+        @factory.callback(:after_create) {}
+        @factory.attributes.should include('after_create callback')
+      end
+
+      it "should raise an InvalidCallbackNameError when defining a callback with an invalid name" do
+        lambda{
+          @factory.callback(:invalid_callback_name) {}
+        }.should raise_error(Factory::InvalidCallbackNameError)
+      end
+    end
+
     describe "after adding an attribute" do
       before do
         @attribute = "attribute"

spec/factory_girl/proxy/build_spec.rb

@@ -45,6 +45,14 @@
     @proxy.result.should == @instance
   end
 
+  it "should run the :after_build callback when retrieving the result" do
+    spy = Object.new
+    stub(spy).foo
+    @proxy.add_callback(:after_build, proc{ spy.foo })
+    @proxy.result
+    spy.should have_received.foo
+  end
+
   describe "when setting an attribute" do
     before do
       stub(@instance).attribute = 'value'

spec/factory_girl/proxy/create_spec.rb

@@ -44,6 +44,12 @@
 
   describe "when asked for the result" do
     before do
+      @build_spy = Object.new
+      @create_spy = Object.new
+      stub(@build_spy).foo
+      stub(@create_spy).foo
+      @proxy.add_callback(:after_build,  proc{ @build_spy.foo })
+      @proxy.add_callback(:after_create, proc{ @create_spy.foo })
       @result = @proxy.result
     end
 
@@ -54,6 +60,11 @@
     it "should return the built instance" do
       @result.should == @instance
     end
+
+    it "should run both the build and the create callbacks" do
+      @build_spy.should have_received.foo
+      @create_spy.should have_received.foo
+    end
   end
 
   describe "when setting an attribute" do

spec/factory_girl/proxy/stub_spec.rb

@@ -45,8 +45,18 @@
       @stub.association(:user).should == @user
     end
 
-    it "should return the actual instance when asked for the result" do
-      @stub.result.should == @instance
+    describe "when asked for the result" do
+      it "should return the actual instance when asked for the result" do
+        @stub.result.should == @instance
+      end
+
+      it "should run the :after_stub callback when asked for the result" do
+        @spy = Object.new
+        stub(@spy).foo
+        @stub.add_callback(:after_stub, proc{ @spy.foo })
+        @stub.result
+        @spy.should have_received.foo
+      end
     end
   end
 

spec/factory_girl/proxy_spec.rb

@@ -25,4 +25,60 @@
   it "should raise an error when asked for the result" do
     lambda { @proxy.result }.should raise_error(NotImplementedError)
   end
+
+  describe "when adding callbacks" do
+    before do
+      @first_block  = proc{ 'block 1' }
+      @second_block = proc{ 'block 2' }
+    end
+    it "should add a callback" do
+      @proxy.add_callback(:after_create, @first_block)
+      @proxy.callbacks[:after_create].should be_eql([@first_block])
+    end
+
+    it "should add multiple callbacks of the same name" do
+      @proxy.add_callback(:after_create, @first_block)
+      @proxy.add_callback(:after_create, @second_block)
+      @proxy.callbacks[:after_create].should be_eql([@first_block, @second_block])
+    end
+
+    it "should add multiple callbacks of different names" do
+      @proxy.add_callback(:after_create, @first_block)
+      @proxy.add_callback(:after_build,  @second_block)
+      @proxy.callbacks[:after_create].should be_eql([@first_block])
+      @proxy.callbacks[:after_build].should be_eql([@second_block])
+    end
+  end
+
+  describe "when running callbacks" do
+    before do
+      @first_spy = Object.new
+      @second_spy = Object.new
+      stub(@first_spy).foo
+      stub(@second_spy).foo
+    end
+
+    it "should run all callbacks with a given name" do
+      @proxy.add_callback(:after_create, proc{ @first_spy.foo })
+      @proxy.add_callback(:after_create, proc{ @second_spy.foo })
+      @proxy.run_callbacks(:after_create)
+      @first_spy.should have_received.foo
+      @second_spy.should have_received.foo
+    end
+
+    it "should only run callbacks with a given name" do
+      @proxy.add_callback(:after_create, proc{ @first_spy.foo })
+      @proxy.add_callback(:after_build,  proc{ @second_spy.foo })
+      @proxy.run_callbacks(:after_create)
+      @first_spy.should have_received.foo
+      @second_spy.should_not have_received.foo
+    end
+
+    it "should pass in the instance if the block takes an argument" do
+      @proxy.instance_variable_set("@instance", @first_spy)
+      @proxy.add_callback(:after_create, proc{|spy| spy.foo })
+      @proxy.run_callbacks(:after_create)
+      @first_spy.should have_received.foo
+    end
+  end
 end

spec/integration_spec.rb

@@ -31,6 +31,17 @@
       f.username  'GuestUser'
     end
 
+    Factory.define :user_with_callbacks, :parent => :user do |f|
+      f.after_stub   {|u| u.first_name = 'Stubby' }
+      f.after_build  {|u| u.first_name = 'Buildy' }
+      f.after_create {|u| u.last_name  = 'Createy' }
+    end
+
+    Factory.define :business do |f|
+      f.name 'Supplier of Awesome'
+      f.association :owner, :factory => :user
+    end
+
     Factory.sequence :email do |n|
       "somebody#{n}@example.com"
     end
@@ -262,4 +273,22 @@
       Factory(:sequence_abuser)
     }.should raise_error(Factory::SequenceAbuseError)
   end
+
+  describe "an instance with callbacks" do
+    it "should run the after_stub callback when stubbing" do
+      @user = Factory.stub(:user_with_callbacks)
+      @user.first_name.should == 'Stubby'
+    end
+
+    it "should run the after_build callback when building" do
+      @user = Factory.build(:user_with_callbacks)
+      @user.first_name.should == 'Buildy'
+    end
+
+    it "should run both the after_build and after_create callbacks when creating" do
+      @user = Factory(:user_with_callbacks)
+      @user.first_name.should == 'Buildy'
+      @user.last_name.should == 'Createy'
+    end
+  end
 end