methods to clear cached classes after reload! is called

Gemfile

@@ -2,7 +2,7 @@ source 'http://rubygems.org'
 
 gemspec
 
-# gem 'neo4j-core', github: 'neo4jrb/neo4j-core', branch: 'master'
+gem 'neo4j-core', github: 'neo4jrb/neo4j-core', branch: 'master'
 # gem 'neo4j-core', path: '../neo4j-core'
 
 # gem 'active_attr', github: 'neo4jrb/active_attr', branch: 'performance'

lib/neo4j/active_node/has_n/association.rb

@@ -20,21 +20,27 @@ def initialize(type, direction, name, options = {type: nil})
         end
 
         def derive_model_class
+          refresh_model_class! if pending_model_refresh?
           return @model_class unless @model_class.nil?
           return nil if relationship_class.nil?
           dir_class = direction == :in ? :from_class : :to_class
           return false if relationship_class.send(dir_class).to_s.to_sym == :any
           relationship_class.send(dir_class)
         end
 
+        def refresh_model_class!
+          @pending_model_refresh = @target_classes_or_nil = nil
+          @model_class = @model_class.name.constantize if @model_class
+        end
+
+        def queue_model_refresh!
+          @pending_model_refresh = true
+        end
+
         def target_class_option(model_class)
           case model_class
           when nil
-            if @target_class_name_from_name
-              "#{association_model_namespace}::#{@target_class_name_from_name}"
-            else
-              @target_class_name_from_name
-            end
+            @target_class_name_from_name ? "#{association_model_namespace}::#{@target_class_name_from_name}" : @target_class_name_from_name
           when Array
             model_class.map { |sub_model_class| target_class_option(sub_model_class) }
           when false
@@ -44,6 +50,10 @@ def target_class_option(model_class)
           end
         end
 
+        def pending_model_refresh?
+          !!@pending_model_refresh
+        end
+
         def target_class_names
           option = target_class_option(derive_model_class)
 
@@ -172,9 +182,7 @@ def validate_init_arguments(type, direction, name, options)
           check_valid_type_and_dir(type, direction)
         end
 
-        VALID_ASSOCIATION_OPTION_KEYS = [:type, :origin, :model_class,
-                                         :rel_class, :dependent, :before,
-                                         :after, :unique]
+        VALID_ASSOCIATION_OPTION_KEYS = [:type, :origin, :model_class, :rel_class, :dependent, :before, :after, :unique]
 
         def validate_association_options!(_association_name, options)
           type_keys = (options.keys & [:type, :origin, :rel_class])

lib/neo4j/active_node/labels.rb

@@ -80,6 +80,12 @@ def self.clear_wrapped_models
       module ClassMethods
         include Neo4j::ActiveNode::QueryMethods
 
+        def before_remove_const
+          associations.each_value(&:queue_model_refresh!)
+          MODELS_FOR_LABELS_CACHE.clear
+          WRAPPED_CLASSES.clear
+        end
+
         # Returns the object with the specified neo4j id.
         # @param [String,Integer] id of node to find
         def find(id)

spec/unit/association_spec.rb

@@ -380,4 +380,47 @@ def self.name
       end
     end
   end
+
+  describe 'model refresh methods' do
+    let(:type) { :has_many }
+    describe '#queue_model_refresh!' do
+      it 'changes the response of #pending_model_refresh' do
+        expect { association.queue_model_refresh! }.to change { association.pending_model_refresh? }
+      end
+    end
+
+    describe 'refresh_model_class!' do
+      context 'with model class set' do
+        before do
+          stub_const('MyModel', Class.new)
+          association.instance_variable_set(:@model_class, MyModel)
+          name = 'MyModel'
+          expect(MyModel).to receive(:name).and_return(name)
+          expect(name).to receive(:constantize).and_return(Class.new)
+        end
+
+        it 'changes the value of #derive_model_class' do
+          expect { association.refresh_model_class! }.to change { association.derive_model_class }
+        end
+
+        it 'resets #pending_model_refresh?' do
+          association.queue_model_refresh!
+          expect { association.refresh_model_class! }.to change { association.pending_model_refresh? }
+        end
+      end
+
+      context 'without model class set' do
+        before { association.instance_variable_set(:@model_class, nil) }
+
+        it 'does not raise an error' do
+          expect { association.refresh_model_class! }.not_to raise_error
+        end
+
+        it 'still resets #pending_model_refresh?' do
+          association.queue_model_refresh!
+          expect { association.refresh_model_class! }.to change { association.pending_model_refresh? }
+        end
+      end
+    end
+  end
 end