model.rb in with_model-2.1.3

- old
+ new

@@ -11,12 +11,12 @@
   # In general, direct use of this class should be avoided. Instead use
   # either the {WithModel high-level API} or {WithModel::Model::DSL low-level API}.
   class Model
     attr_writer :model_block, :table_block, :table_options
 
-    # @param name The constant name (as a symbol) to assign the model class to.
-    # @param superclass The superclass for the created class. Should
+    # @param [Symbol] name The constant name to assign the model class to.
+    # @param [Class] superclass The superclass for the created class. Should
     #   have `ActiveRecord::Base` as an ancestor.
     def initialize(name, superclass: ActiveRecord::Base)
       @name = name.to_sym
       @model_block = nil
       @table_block = nil
@@ -33,11 +33,11 @@
       setup_model
     end
 
     def destroy
       stubber.unstub_const
-      remove_from_superclass_descendants
+      cleanup_descendants_tracking
       reset_dependencies_cache
       table.destroy
       @model = nil
     end
 
@@ -51,12 +51,15 @@
       @model.table_name = table_name
       @model.class_eval(&@model_block) if @model_block
       @model.reset_column_information
     end
 
-    def remove_from_superclass_descendants
-      return unless @model.superclass.respond_to?(:direct_descendants)
-      @model.superclass.direct_descendants.delete(@model)
+    def cleanup_descendants_tracking
+      if defined?(ActiveSupport::DescendantsTracker)
+        ActiveSupport::DescendantsTracker.class_variable_get(:@@direct_descendants).delete(ActiveRecord::Base)
+      elsif @model.superclass.respond_to?(:direct_descendants)
+        @model.superclass.direct_descendants.delete(@model)
+      end
     end
 
     def reset_dependencies_cache
       return unless defined?(ActiveSupport::Dependencies::Reference)
       ActiveSupport::Dependencies::Reference.clear!