Support cache_classes: true descendants filtering

When cache_classes is true, we can no longer rely on the Active Support
dependencies tracker, so we use our own separate tracker to filter out
destroyed with_model models from Class#descendants and Class#subclasses.

Fixes #35

Author: nertzy

.rubocop.yml

@@ -29,5 +29,8 @@ RSpec/Be:
 RSpec/BeforeAfterAll:
   Enabled: false
 
+RSpec/ExpectInHook:
+  Enabled: false
+
 Style/Documentation:
   Enabled: false

Gemfile

@@ -17,6 +17,7 @@ end
 gem "bigdecimal"
 gem "bundler"
 gem "debug"
+gem "logger"
 gem "minitest"
 gem "mutex_m"
 gem "rake"

lib/with_model/descendants_tracker.rb

@@ -0,0 +1,86 @@
+require "active_support/descendants_tracker"
+
+module WithModel
+  module DescendantsTracker
+    if RUBY_ENGINE == "ruby"
+      # On MRI `ObjectSpace::WeakMap` keys are weak references.
+      # So we can simply use WeakMap as a `Set`.
+      class WeakSet < ObjectSpace::WeakMap # :nodoc:
+        alias_method :to_a, :keys
+
+        def <<(object)
+          self[object] = true
+        end
+      end
+    else
+      # On TruffleRuby `ObjectSpace::WeakMap` keys are strong references.
+      # So we use `object_id` as a key and the actual object as a value.
+      #
+      # JRuby for now doesn't have Class#descendant, but when it will, it will likely
+      # have the same WeakMap semantic than Truffle so we future proof this as much as possible.
+      class WeakSet # :nodoc:
+        def initialize
+          @map = ObjectSpace::WeakMap.new
+        end
+
+        def [](object)
+          @map.key?(object.object_id)
+        end
+        alias_method :include?, :[]
+
+        def []=(object, _present)
+          @map[object.object_id] = object
+        end
+
+        def to_a
+          @map.values
+        end
+
+        def <<(object)
+          self[object] = true
+        end
+      end
+    end
+    @excluded_descendants = WeakSet.new
+
+    class << self
+      def clear(classes) # :nodoc:
+        classes.each do |klass|
+          @excluded_descendants << klass
+          klass.descendants.each do |descendant|
+            @excluded_descendants << descendant
+          end
+        end
+      end
+
+      def reject!(classes) # :nodoc:
+        if @excluded_descendants
+          classes.reject! { |d| @excluded_descendants.include?(d) }
+        end
+        classes
+      end
+    end
+
+    module ReloadedClassesFiltering # :nodoc:
+      def subclasses
+        WithModel::DescendantsTracker.reject!(super)
+      end
+
+      def descendants
+        WithModel::DescendantsTracker.reject!(super)
+      end
+    end
+  end
+end
+
+class Class
+  prepend WithModel::DescendantsTracker::ReloadedClassesFiltering
+end
+
+module ActiveSupport
+  module DescendantsTracker
+    class << self
+      attr_reader :clear_disabled
+    end
+  end
+end

lib/with_model/model.rb

@@ -1,9 +1,11 @@
 # frozen_string_literal: true
 
+require "logger"
 require "active_record"
 require "active_support/core_ext/string/inflections"
 require "English"
 require "with_model/constant_stubber"
+require "with_model/descendants_tracker"
 require "with_model/methods"
 require "with_model/table"
 
@@ -38,6 +40,7 @@ def destroy
       cleanup_descendants_tracking
       reset_dependencies_cache
       table.destroy
+      WithModel::DescendantsTracker.clear([@model])
       @model = nil
     end
 
@@ -54,15 +57,8 @@ def setup_model
     end
 
     def cleanup_descendants_tracking
-      if defined?(ActiveSupport::DescendantsTracker)
-        if ActiveSupport::VERSION::MAJOR >= 7
-          ActiveSupport::DescendantsTracker.clear([@model])
-        else
-          ActiveSupport::DescendantsTracker.class_variable_get(:@@direct_descendants).delete(ActiveRecord::Base)
-        end
-      elsif @model.superclass.respond_to?(:direct_descendants)
-        @model.superclass.direct_descendants.delete(@model)
-      end
+      ActiveSupport::DescendantsTracker.clear([@model]) \
+        unless ActiveSupport::DescendantsTracker.clear_disabled
     end
 
     def reset_dependencies_cache

spec/descendants_tracking_spec.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+describe "Descendants tracking" do # rubocop:disable RSpec/DescribeClass
+  with_model :BlogPost do
+    model do
+      def self.inspect
+        "BlogPost class #{object_id}"
+      end
+    end
+  end
+
+  def blog_post_classes
+    ActiveRecord::Base.descendants.select do |c|
+      c.table_name == BlogPost.table_name
+    end
+  end
+
+  shared_examples "clearing descendants between test runs" do
+    it "includes the correct model class in descendants on the first test run" do
+      expect(blog_post_classes).to eq [BlogPost]
+    end
+
+    it "includes the correct model class in descendants on the second test run" do
+      expect(blog_post_classes).to eq [BlogPost]
+    end
+  end
+
+  context "with ActiveSupport::DescendantsTracker (cache_classes: true)" do
+    before do
+      expect(ActiveSupport::DescendantsTracker.clear_disabled).to be_falsey
+      expect { ActiveSupport::DescendantsTracker.clear([]) }.not_to raise_exception
+    end
+
+    include_examples "clearing descendants between test runs"
+  end
+
+  context "without ActiveSupport::DescendantsTracker (cache_classes: false)" do
+    before do
+      ActiveSupport::DescendantsTracker.disable_clear!
+      expect(ActiveSupport::DescendantsTracker.clear_disabled).to be_truthy
+      expect { ActiveSupport::DescendantsTracker.clear([]) }.to raise_exception(RuntimeError)
+    end
+
+    include_examples "clearing descendants between test runs"
+  end
+end

spec/with_model_spec.rb

@@ -320,30 +320,6 @@ def my_method
     end
   end
 
-  context "with ActiveSupport::DescendantsTracker" do
-    with_model :BlogPost do
-      model do
-        def self.inspect
-          "BlogPost class #{object_id}"
-        end
-      end
-    end
-
-    def blog_post_classes
-      ActiveRecord::Base.descendants.select do |c|
-        c.table_name == BlogPost.table_name
-      end
-    end
-
-    it "includes the correct model class in descendants on the first test run" do
-      expect(blog_post_classes).to eq [BlogPost]
-    end
-
-    it "includes the correct model class in descendants on the second test run" do
-      expect(blog_post_classes).to eq [BlogPost]
-    end
-  end
-
   context "with_model can be run within RSpec :all hook" do
     with_model :BlogPost, scope: :all do
       table do |t|
@@ -396,10 +372,6 @@ class ApplicationRecordInDifferentDatabase < ActiveRecord::Base # standard:disab
       establish_connection(ActiveRecord::Base.connection_pool.db_config.configuration_hash)
     end
 
-    after(:all) do
-      Object.__send__(:remove_const, "ApplicationRecordInDifferentDatabase")
-    end
-
     with_model :BlogPost, superclass: ApplicationRecordInDifferentDatabase do
       table do |t|
         t.string "title"