Allow extension from within Associated Object

lib/active_record/associated_object.rb

@@ -15,6 +15,10 @@ def inherited(klass)
       klass.delegate :record_klass, :attribute_name, to: :class
     end
 
+    def extension(&block)
+      record_klass.class_eval(&block)
+    end
+
     def respond_to_missing?(...) = record_klass.respond_to?(...) || super
     delegate :unscoped, :transaction, :primary_key, to: :record_klass
 

test/active_record/associated_object_test.rb

@@ -130,4 +130,9 @@ def test_active_job_integration
   def test_calling_method
     assert @rating.great?
   end
+
+  def test_record_klass_extension
+    assert_predicate Post::Comment.great.first, :rated_great?
+    assert_match /test\/boot\/associated_object/, Post::Comment.instance_method(:rated_great?).source_location.first
+  end
 end

test/boot/associated_object.rb

@@ -41,6 +41,12 @@ def publish
 Post.has_object :publisher, after_update_commit: true, before_destroy: :prevent_errant_post_destroy
 
 class Post::Comment::Rating < ActiveRecord::AssociatedObject
+  extension do
+    scope :great, -> { where(body: "First!!!!") }
+
+    def rated_great? = rating.great?
+  end
+
   kredis_flag :moderated
 
   def great?