Fixed reverse delete rule with inheritance (#197)

mongoengine/document.py

@@ -7,7 +7,7 @@ from bson.dbref import DBRef
 from mongoengine import signals, queryset
 
 from base import (DocumentMetaclass, TopLevelDocumentMetaclass, BaseDocument,
-                  BaseDict, BaseList, ALLOW_INHERITANCE)
+                  BaseDict, BaseList, ALLOW_INHERITANCE, get_document)
 from queryset import OperationError, NotUniqueError
 from connection import get_db, DEFAULT_CONNECTION_NAME
 
@@ -421,6 +421,15 @@ class Document(BaseDocument):
         """This method registers the delete rules to apply when removing this
         object.
         """
+        classes = [get_document(class_name)
+                    for class_name in cls._subclasses
+                    if class_name != cls.__name__] + [cls]
+        documents = [get_document(class_name)
+                     for class_name in document_cls._subclasses
+                     if class_name != document_cls.__name__] + [document_cls]
+
+        for cls in classes:
+            for document_cls in documents:
                 delete_rules = cls._meta.get('delete_rules') or {}
                 delete_rules[(document_cls, field_name)] = rule
                 cls._meta['delete_rules'] = delete_rules

tests/document/instance.py

@@ -1378,6 +1378,42 @@ class InstanceTest(unittest.TestCase):
         author.delete()
         self.assertEqual(len(BlogPost.objects), 0)
 
+    def test_reverse_delete_rule_with_document_inheritance(self):
+        """Ensure that a referenced document is also deleted upon deletion
+        of a child document.
+        """
+
+        class Writer(self.Person):
+            pass
+
+        class BlogPost(Document):
+            content = StringField()
+            author = ReferenceField(self.Person, reverse_delete_rule=CASCADE)
+            reviewer = ReferenceField(self.Person, reverse_delete_rule=NULLIFY)
+
+        self.Person.drop_collection()
+        BlogPost.drop_collection()
+
+        author = Writer(name='Test User')
+        author.save()
+
+        reviewer = Writer(name='Re Viewer')
+        reviewer.save()
+
+        post = BlogPost(content='Watched some TV')
+        post.author = author
+        post.reviewer = reviewer
+        post.save()
+
+        reviewer.delete()
+        self.assertEqual(len(BlogPost.objects), 1)
+        self.assertEqual(BlogPost.objects.get().reviewer, None)
+
+        # Delete the Writer should lead to deletion of the BlogPost
+        author.delete()
+        self.assertEqual(len(BlogPost.objects), 0)
+
+
     def test_reverse_delete_rule_cascade_and_nullify_complex_field(self):
         """Ensure that a referenced document is also deleted upon deletion for
         complex fields.