lib/mongomapper/associations/belongs_to_proxy.rb in jnunemaker-mongomapper-0.2.0 vs lib/mongomapper/associations/belongs_to_proxy.rb in jnunemaker-mongomapper-0.3.0

@@ -1,26 +1,22 @@
 module MongoMapper
   module Associations
     class BelongsToProxy < Proxy
-      def replace(v)
-        ref_id = "#{@association.name}_id"
-
-        if v
-          v.save if v.new?
-          @owner.__send__(:write_attribute, ref_id, v.id)
-        else
-          @owner.__send__(:write_attribute, ref_id, nil)
+      def replace(doc)
+        if doc
+          doc.save if doc.new?
+          id = doc.id
         end
-
+        
+        @owner.send("#{@association.belongs_to_key_name}=", id)
         reload_target
       end
 
       protected
-      def find_target
-        ref = @owner.__send__(:read_attribute, "#{@association.name}_id")
-        if ref
-          @association.klass.find(ref)
+        def find_target
+          if ref = @owner.send(@association.belongs_to_key_name)
+            @association.klass.find(ref)
+          end
         end
-      end
     end
   end
 end