From 445f9453c4d305c394bb833dcd16ddce4b3ab2b8 Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Tue, 22 Jan 2013 16:38:07 +0000
Subject: [PATCH] Fixed reverse delete rule with inheritance (#197)

---
 mongoengine/document.py    | 17 +++++++++++++----
 tests/document/instance.py | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/mongoengine/document.py b/mongoengine/document.py
index fff7efad..f40f1c9f 100644
--- a/mongoengine/document.py
+++ b/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,9 +421,18 @@ class Document(BaseDocument):
         """This method registers the delete rules to apply when removing this
         object.
         """
-        delete_rules = cls._meta.get('delete_rules') or {}
-        delete_rules[(document_cls, field_name)] = rule
-        cls._meta['delete_rules'] = delete_rules
+        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
 
     @classmethod
     def drop_collection(cls):
diff --git a/tests/document/instance.py b/tests/document/instance.py
index c8a1b110..07c4f0e0 100644
--- a/tests/document/instance.py
+++ b/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.