From b7111b2ce33a4d125a9cec2daa4cfa980063e867 Mon Sep 17 00:00:00 2001
From: Jon Leighton <j@jonathanleighton.com>
Date: Sun, 7 Sep 2008 14:42:33 +0100
Subject: [PATCH] Support for updating a belongs to association from the foreign key (without saving and reloading the record)

---
 activerecord/lib/active_record/associations.rb     |    9 ++++++++-
 .../associations/belongs_to_associations_test.rb   |   13 +++++++++++++
 activerecord/test/cases/associations_test.rb       |   12 ++++++++++++
 3 files changed, 33 insertions(+), 1 deletions(-)

diff --git a/activerecord/lib/active_record/associations.rb b/activerecord/lib/active_record/associations.rb
index 6405071..8844c50 100755
--- a/activerecord/lib/active_record/associations.rb
+++ b/activerecord/lib/active_record/associations.rb
@@ -1246,7 +1246,7 @@ module ActiveRecord
 
             association = instance_variable_get(ivar) if instance_variable_defined?(ivar)
 
-            if association.nil? || force_reload
+            if association.nil? || !association.loaded? || force_reload
               association = association_proxy_class.new(self, reflection)
               retval = association.reload
               if retval.nil? and association_proxy_class == BelongsToAssociation
@@ -1276,6 +1276,13 @@ module ActiveRecord
               instance_variable_set(ivar, new_value.nil? ? nil : association)
             end
           end
+          
+          if association_proxy_class == BelongsToAssociation
+            define_method("#{reflection.primary_key_name}=") do |target_id|
+              instance_variable_get(ivar).reset if instance_variable_defined?(ivar)
+              write_attribute(reflection.primary_key_name, target_id)
+            end
+          end
 
           define_method("set_#{reflection.name}_target") do |target|
             return if target.nil? and association_proxy_class == BelongsToAssociation
diff --git a/activerecord/test/cases/associations/belongs_to_associations_test.rb b/activerecord/test/cases/associations/belongs_to_associations_test.rb
index 9c718c4..37aaa0b 100644
--- a/activerecord/test/cases/associations/belongs_to_associations_test.rb
+++ b/activerecord/test/cases/associations/belongs_to_associations_test.rb
@@ -46,6 +46,19 @@ class BelongsToAssociationsTest < ActiveRecord::TestCase
     citibank.firm = apple
     assert_equal apple.id, citibank.firm_id
   end
+  
+  def test_foreign_key_assignment
+    # Test using an existing record
+    signals37 = accounts(:signals37)
+    assert_equal companies(:first_firm), signals37.firm
+    signals37.firm_id = companies(:another_firm).id
+    assert_equal companies(:another_firm), signals37.firm
+    
+    # Test using a new record
+    account = Account.new
+    account.firm_id = companies(:another_firm).id
+    assert_equal companies(:another_firm), account.firm
+  end
 
   def test_no_unexpected_aliasing
     first_firm = companies(:first_firm)
diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb
index 0b2731e..7cecdc3 100644
--- a/activerecord/test/cases/associations_test.rb
+++ b/activerecord/test/cases/associations_test.rb
@@ -181,6 +181,18 @@ class AssociationProxyTest < ActiveRecord::TestCase
     p = setup_dangling_association
     assert_nil p.author.reset
   end
+  
+  def test_reset_loads_association_next_time
+    welcome = posts(:welcome)
+    david = authors(:david)
+    author_assoc = welcome.author
+    
+    assert_equal david, welcome.author # So we can be sure the test works correctly
+    author_assoc.reset
+    assert !author_assoc.loaded?
+    assert_nil author_assoc.target
+    assert_equal david, welcome.author
+  end
 
   def test_reload_returns_assocition
     david = developers(:david)
-- 
1.5.4.3