From 1631788ab6bfe66779f226f750b9f235c1f6c7ea Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Wed, 24 Aug 2011 13:37:39 +0100
Subject: [PATCH] Now Raise an exception if subclasses are missing at
 querytime.

Beats returning None thanks to #aid for mentioning it on IRC
---
 mongoengine/base.py |  7 +++++--
 tests/document.py   | 30 +++++++++++++++++++++++++++---
 tests/fields.py     |  3 ++-
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/mongoengine/base.py b/mongoengine/base.py
index 6be5c3de..8d0c470b 100644
--- a/mongoengine/base.py
+++ b/mongoengine/base.py
@@ -671,7 +671,6 @@ class BaseDocument(object):
         # get the class name from the document, falling back to the given
         # class if unavailable
         class_name = son.get(u'_cls', cls._class_name)
-
         data = dict((str(key), value) for key, value in son.items())
 
         if '_types' in data:
@@ -686,7 +685,11 @@ class BaseDocument(object):
             if class_name not in subclasses:
                 # Type of document is probably more generic than the class
                 # that has been queried to return this SON
-                return None
+                raise NotRegistered("""
+                        `%s` has not been registered in the document registry.
+                        Importing the document class automatically registers it,
+                        has it been imported?
+                    """.strip() % class_name)
             cls = subclasses[class_name]
 
         present_fields = data.keys()
diff --git a/tests/document.py b/tests/document.py
index d6d70289..b76b6f92 100644
--- a/tests/document.py
+++ b/tests/document.py
@@ -12,7 +12,7 @@ import weakref
 from fixtures import Base, Mixin, PickleEmbedded, PickleTest
 
 from mongoengine import *
-from mongoengine.base import BaseField
+from mongoengine.base import _document_registry, NotRegistered
 from mongoengine.connection import _get_db
 
 
@@ -740,7 +740,6 @@ class DocumentTest(unittest.TestCase):
         post1.save()
         BlogPost.drop_collection()
 
-
     def test_geo_indexes_recursion(self):
 
         class User(Document):
@@ -799,7 +798,6 @@ class DocumentTest(unittest.TestCase):
         post2 = BlogPost(title='test2', slug='test')
         self.assertRaises(OperationError, post2.save)
 
-
     def test_unique_with(self):
         """Ensure that unique_with constraints are applied to fields.
         """
@@ -978,6 +976,32 @@ class DocumentTest(unittest.TestCase):
 
         User.drop_collection()
 
+
+    def test_document_not_registered(self):
+
+        class Place(Document):
+            name = StringField()
+
+        class NicePlace(Place):
+            pass
+
+        Place.drop_collection()
+
+        Place(name="London").save()
+        NicePlace(name="Buckingham Palace").save()
+
+        # Mimic Place and NicePlace definitions being in a different file
+        # and the NicePlace model not being imported in at query time.
+        @classmethod
+        def _get_subclasses(cls):
+            return {}
+        Place._get_subclasses = _get_subclasses
+
+        def query_without_importing_nice_place():
+            print Place.objects.all()
+        self.assertRaises(NotRegistered, query_without_importing_nice_place)
+
+
     def test_creation(self):
         """Ensure that document may be created using keyword arguments.
         """
diff --git a/tests/fields.py b/tests/fields.py
index f8aeb86c..f9734900 100644
--- a/tests/fields.py
+++ b/tests/fields.py
@@ -1062,6 +1062,7 @@ class FieldTest(unittest.TestCase):
         Post.drop_collection()
         User.drop_collection()
 
+
     def test_generic_reference_document_not_registered(self):
         """Ensure dereferencing out of the document registry throws a
         `NotRegistered` error.
@@ -1445,7 +1446,7 @@ class FieldTest(unittest.TestCase):
         a.save()
 
         self.assertEqual(a.counter, 2)
-        
+
         a = Animal.objects.first()
         self.assertEqual(a.counter, 2)
         a.reload()