fix association :extend option

2 files changed, 17 insertions, 6 deletions

diff --git a/activerecord/lib/active_record/associations/builder/association.rb b/activerecord/lib/active_record/associations/builder/association.rb
index 331b724814..4a78a9c076 100644
--- a/activerecord/lib/active_record/associations/builder/association.rb
+++ b/activerecord/lib/active_record/associations/builder/association.rb
@@ -23,6 +23,11 @@ module ActiveRecord::Associations::Builder
         @scope   = nil
         @options = scope
       end
+
+      if @scope && @scope.arity == 0
+        prev_scope = @scope
+        @scope = proc { instance_exec(&prev_scope) }
+      end
     end
 
     def mixin
diff --git a/activerecord/lib/active_record/associations/builder/collection_association.rb b/activerecord/lib/active_record/associations/builder/collection_association.rb
index 98ca94b42a..af81af4ad2 100644
--- a/activerecord/lib/active_record/associations/builder/collection_association.rb
+++ b/activerecord/lib/active_record/associations/builder/collection_association.rb
@@ -6,7 +6,7 @@ module ActiveRecord::Associations::Builder
       super + [:table_name, :finder_sql, :counter_sql, :before_add, :after_add, :before_remove, :after_remove]
     end
 
-    attr_reader :block_extension
+    attr_reader :block_extension, :extension_module
 
     def initialize(*args, &extension)
       super(*args)
@@ -27,18 +27,24 @@ module ActiveRecord::Associations::Builder
     private
 
       def wrap_block_extension
-        options[:extend] = Array(options[:extend])
-
         if block_extension
+          @extension_module = mod = Module.new(&block_extension)
           silence_warnings do
-            model.parent.const_set(extension_module_name, Module.new(&block_extension))
+            model.parent.const_set(extension_module_name, mod)
+          end
+
+          prev_scope = @scope
+
+          if prev_scope
+            @scope = proc { |owner| instance_exec(owner, &prev_scope).extending(mod) }
+          else
+            @scope = proc { extending(mod) }
           end
-          options[:extend].push("#{model.parent}::#{extension_module_name}".constantize)
         end
       end
 
       def extension_module_name
-        @extension_module_name ||= "#{model.to_s.demodulize}#{name.to_s.camelize}AssociationExtension"
+        @extension_module_name ||= "#{model.name.demodulize}#{name.to_s.camelize}AssociationExtension"
       end
 
       def define_callback(callback_name)