lib/trestle/registry.rb in trestle-0.10.0 vs lib/trestle/registry.rb in trestle-0.10.1

@@ -13,49 +13,52 @@
       @admins.values.sort_by(&:admin_name).each(&block)
     end
 
     def reset!
       @admins = {}
-      @models = {}
+      @models = nil
     end
 
     def empty?
       none?
     end
 
-    def register(admin, register_model: true)
+    def register(admin)
       @admins[admin.admin_name] = admin
-
-      if register_model && register_admin_for_model_loookup?(admin)
-        @models[admin.model.name] ||= admin
-      end
-
-      admin
     end
 
     def lookup_admin(admin)
       # Given object is already an admin class
       return admin if admin.is_a?(Class) && admin < Trestle::Admin
 
+      # Given object is already an admin instance
+      return admin if admin.is_a?(Trestle::Admin)
+
       @admins[admin.to_s]
     end
     alias lookup lookup_admin
 
     def lookup_model(model)
       # Lookup each class in the model's ancestor chain
       while model
-        admin = @models[model.name]
+        admin = models[model.name]
         return admin if admin
 
         model = model.superclass
       end
 
       # No admin found
       nil
     end
 
   private
-    def register_admin_for_model_loookup?(admin)
-      admin.respond_to?(:model) && !(admin.respond_to?(:singular?) && admin.singular?)
+    def models
+      @models ||= @admins.values.inject({}) { |result, admin|
+        if admin.respond_to?(:register_model?) && admin.register_model?
+          result[admin.model.name] ||= admin
+        end
+
+        result
+      }
     end
   end
 end