lib/active_null.rb in active_null-0.0.5 vs lib/active_null.rb in active_null-0.0.6

@@ -1,18 +1,17 @@
 require 'active_null/version'
 require 'active_null/null_model_builder'
 
 module ActiveNull
   def null
-    @null_class.get
+    null_class.get
   end
 
   def null_model(&block)
-    null_model_overrides = if block_given?
+    @null_model_overrides = if block_given?
       Module.new.tap { |m| m.module_eval(&block) }
     end
-    @null_class = NullModelBuilder.new(self, null_model_overrides).build
   end
 
   def find_by(*args, &block)
     super || null
   end
@@ -27,9 +26,13 @@
             #{relation.klass.name}.null
           end
         CODE
       end
     end
+  end
+
+  def null_class
+    @null_class ||= NullModelBuilder.new(self, @null_model_overrides).build
   end
 
   def self.extended(klass)
     klass.class_eval <<-CODE
       after_initialize do