From d3a5d854857856afc85944e38b2a4e1f46df8a26 Mon Sep 17 00:00:00 2001
From: Joe Ferris <jferris@thoughtbot.com>
Date: Wed, 28 May 2008 20:00:46 -0400
Subject: [PATCH] Implemented add_attribute and attributes for Factory

---
 lib/factory_girl/factory.rb | 47 ++++++++++++++++++++++++++
 test/factory_test.rb        | 67 +++++++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+)

diff --git a/lib/factory_girl/factory.rb b/lib/factory_girl/factory.rb
index cd9dfd0..299a80c 100644
--- a/lib/factory_girl/factory.rb
+++ b/lib/factory_girl/factory.rb
@@ -23,6 +23,53 @@ class Factory
 
   def initialize (name) #:nodoc:
     @name = name
+    @static_attributes = {}
+    @lazy_attributes = {}
+  end
+
+  # Adds an attribute that should be assigned or stubbed on generated instances for this factory.
+  #
+  # This method should be called with either a value or block, but not both. If
+  # called with a block, the attribute will be generated "lazily," whenever an
+  # instance is generated. Lazy attribute blocks will not be called if that
+  # attribute is overriden for a specific instance.
+  #
+  # Arguments:
+  #   name: (Symbol)
+  #     The name of this attribute. This will be assigned using :"#{name}=" for
+  #     generated instances, and stubbed out with this name for generated
+  #     stubs.
+  #   value: (Object)
+  #     If no block is given, this value will be used for this attribute.
+  def add_attribute (name, value = nil, &block)
+    if block_given?
+      unless value.nil?
+        raise ArgumentError, "Both value and block given"
+      end
+      @lazy_attributes[name] = block
+    else
+      @static_attributes[name] = value
+    end
+  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.
+  #
+  # Arguments:
+  #   attrs: (Hash)
+  #     Attributes to overwrite for this set.
+  #
+  # Returns:
+  #   A set of attributes that can be used to build an instance of the class
+  #   this factory generates. (Hash)
+  def attributes (attrs = {})
+    result = {}
+    @lazy_attributes.each do |name, block|
+      result[name] = block.call unless attrs.key?(name)
+    end
+    result.update(@static_attributes)
+    result.update(attrs)
   end
 
 end
diff --git a/test/factory_test.rb b/test/factory_test.rb
index 7eb9641..fa86f1e 100644
--- a/test/factory_test.rb
+++ b/test/factory_test.rb
@@ -42,6 +42,73 @@ class FactoryTest < Test::Unit::TestCase
       assert_equal @name, @factory.name
     end
 
+    context "when adding an attribute with a value parameter" do
+
+      setup do
+        @attr  = :name
+        @value = 'Elvis lives!'
+        @factory.add_attribute(@attr, @value)
+      end
+
+      should "include that value in the generated attributes hash" do
+        assert_equal @value, @factory.attributes[@attr]
+      end
+
+    end
+
+    context "when adding an attribute with a block" do
+
+      setup do
+        @attr = :name
+      end
+
+      should "not evaluate the block when the attribute is added" do
+        @factory.add_attribute(@attr) { flunk }
+      end
+
+      should "evaluate the block when attributes are generated" do
+        called = false
+        @factory.add_attribute(@attr) do
+          called = true
+        end
+        @factory.attributes
+        assert called
+      end
+
+      should "use the result of the block as the value of the attribute" do
+        value = "Watch out for snakes!"
+        @factory.add_attribute(@attr) { value }
+        assert_equal value, @factory.attributes[@attr]
+      end
+
+    end
+
+    should "not allow attributes to be added with both a value parameter and a block" do
+      assert_raise(ArgumentError) do
+        @factory.add_attribute(:name, 'value') {}
+      end
+    end
+
+    context "when overriding generated attributes with a hash" do
+
+      setup do
+        @attr  = :name
+        @value = 'The price is right!'
+        @hash  = { @attr => @value }
+      end
+
+      should "return the overridden value in the generated attributes" do
+        @factory.add_attribute(@attr, 'The price is wrong, Bob!')
+        assert_equal @value, @factory.attributes(@hash)[@attr]
+      end
+
+      should "not call a lazy attribute block for an overridden attribute" do
+        @factory.add_attribute(@attr) { flunk }
+        @factory.attributes(@hash)
+      end
+
+    end
+
   end
 
 end