Add block syntax to HasManyAssociation#build. [#502 state:resolve]

Signed-off-by: Pratik Naik <pratiknaik@gmail.com>

2 files changed, 37 insertions, 3 deletions

diff --git a/activerecord/lib/active_record/associations/association_collection.rb b/activerecord/lib/active_record/associations/association_collection.rb
index eb39714909..04be59e89d 100644
--- a/activerecord/lib/active_record/associations/association_collection.rb
+++ b/activerecord/lib/active_record/associations/association_collection.rb
@@ -78,11 +78,14 @@ module ActiveRecord
         @loaded = false
       end
 
-      def build(attributes = {})
+      def build(attributes = {}, &block)
         if attributes.is_a?(Array)
-          attributes.collect { |attr| build(attr) }
+          attributes.collect { |attr| build(attr, &block) }
         else
-          build_record(attributes) { |record| set_belongs_to_association_for(record) }
+          build_record(attributes) do |record|
+            block.call(record) if block_given?
+            set_belongs_to_association_for(record)
+          end
         end
       end
 
diff --git a/activerecord/test/cases/associations/has_many_associations_test.rb b/activerecord/test/cases/associations/has_many_associations_test.rb
index e90edbb213..b9c7ec6377 100644
--- a/activerecord/test/cases/associations/has_many_associations_test.rb
+++ b/activerecord/test/cases/associations/has_many_associations_test.rb
@@ -425,6 +425,37 @@ class HasManyAssociationsTest < ActiveRecord::TestCase
     assert_equal 2, first_topic.replies.to_ary.size
   end
 
+  def test_build_via_block
+    company = companies(:first_firm)
+    new_client = assert_no_queries { company.clients_of_firm.build {|client| client.name = "Another Client" } }
+    assert !company.clients_of_firm.loaded?
+
+    assert_equal "Another Client", new_client.name
+    assert new_client.new_record?
+    assert_equal new_client, company.clients_of_firm.last
+    company.name += '-changed'
+    assert_queries(2) { assert company.save }
+    assert !new_client.new_record?
+    assert_equal 2, company.clients_of_firm(true).size
+  end
+
+  def test_build_many_via_block
+    company = companies(:first_firm)
+    new_clients = assert_no_queries do
+      company.clients_of_firm.build([{"name" => "Another Client"}, {"name" => "Another Client II"}]) do |client|
+        client.name = "changed"
+      end
+    end
+
+    assert_equal 2, new_clients.size
+    assert_equal "changed", new_clients.first.name
+    assert_equal "changed", new_clients.last.name
+
+    company.name += '-changed'
+    assert_queries(3) { assert company.save }
+    assert_equal 3, company.clients_of_firm(true).size
+  end
+
   def test_create_without_loading_association
     first_firm  = companies(:first_firm)
     Firm.column_names