Fix Django's QuerySet.delete() behavior for models that use "invalid"

Signed-off-by: James Ren <[email protected]>

--- autotest/frontend/afe/model_logic.py        2010-02-26 15:28:08.000000000 
-0800
+++ autotest/frontend/afe/model_logic.py        2010-02-26 15:28:08.000000000 
-0800
@@ -530,7 +530,24 @@
             getattr(base_object, related_list_name).append(related_object)
 
 
-class ValidObjectsManager(ExtendedManager):
+class ModelWithInvalidQuerySet(dbmodels.query.QuerySet):
+    """
+    QuerySet that handles delete() properly for models with an "invalid" bit
+    """
+    def delete(self):
+        for model in self:
+            model.delete()
+
+
+class ModelWithInvalidManager(ExtendedManager):
+    """
+    Manager for objects with an "invalid" bit
+    """
+    def get_query_set(self):
+        return ModelWithInvalidQuerySet(self.model)
+
+
+class ValidObjectsManager(ModelWithInvalidManager):
     """
     Manager returning only objects with invalid=False.
     """
--- autotest/frontend/afe/models.py     2010-02-26 15:28:08.000000000 -0800
+++ autotest/frontend/afe/models.py     2010-02-26 15:28:08.000000000 -0800
@@ -47,7 +47,7 @@
                                   editable=settings.FULL_ADMIN)
 
     name_field = 'name'
-    objects = model_logic.ExtendedManager()
+    objects = model_logic.ModelWithInvalidManager()
     valid_objects = model_logic.ValidObjectsManager()
 
 
@@ -91,7 +91,7 @@
     only_if_needed = dbmodels.BooleanField(default=False)
 
     name_field = 'name'
-    objects = model_logic.ExtendedManager()
+    objects = model_logic.ModelWithInvalidManager()
     valid_objects = model_logic.ValidObjectsManager()
     atomic_group = dbmodels.ForeignKey(AtomicGroup, null=True, blank=True)
 
@@ -221,7 +221,7 @@
     dirty = dbmodels.BooleanField(default=True, editable=settings.FULL_ADMIN)
 
     name_field = 'hostname'
-    objects = model_logic.ExtendedManager()
+    objects = model_logic.ModelWithInvalidManager()
     valid_objects = model_logic.ValidObjectsManager()
 
 
--- autotest/frontend/afe/models_test.py        2010-02-26 15:28:08.000000000 
-0800
+++ autotest/frontend/afe/models_test.py        2010-02-26 15:28:08.000000000 
-0800
@@ -141,5 +141,77 @@
         self.assertEquals(entry.execution_path(), '1-autotest_system/subdir')
 
 
+class ModelWithInvalidTest(unittest.TestCase,
+                           frontend_test_utils.FrontendTestMixin):
+    def setUp(self):
+        self._frontend_common_setup()
+
+
+    def tearDown(self):
+        self._frontend_common_teardown()
+
+
+    def test_model_with_invalid_delete(self):
+        self.assertFalse(self.hosts[0].invalid)
+        self.hosts[0].delete()
+        self.assertTrue(self.hosts[0].invalid)
+        self.assertTrue(models.Host.objects.get(id=self.hosts[0].id))
+
+
+    def test_model_with_invalid_delete_queryset(self):
+        for host in self.hosts:
+            self.assertFalse(host.invalid)
+
+        hosts = models.Host.objects.all()
+        hosts.delete()
+        self.assertEqual(hosts.count(), len(self.hosts))
+
+        for host in hosts:
+            self.assertTrue(host.invalid)
+
+
+    def test_cloned_queryset_delete(self):
+        """
+        Make sure that a cloned queryset maintains the custom delete()
+        """
+        to_delete = ('host1', 'host2')
+
+        for host in self.hosts:
+            self.assertFalse(host.invalid)
+
+        hosts = models.Host.objects.all().filter(hostname__in=to_delete)
+        hosts.delete()
+        all_hosts = models.Host.objects.all()
+        self.assertEqual(all_hosts.count(), len(self.hosts))
+
+        for host in all_hosts:
+            if host.hostname in to_delete:
+                self.assertTrue(
+                        host.invalid,
+                        '%s.invalid expected to be True' % host.hostname)
+            else:
+                self.assertFalse(
+                        host.invalid,
+                        '%s.invalid expected to be False' % host.hostname)
+
+
+    def test_normal_delete(self):
+        job = self._create_job(hosts=[1])
+        self.assertEqual(1, models.Job.objects.all().count())
+
+        job.delete()
+        self.assertEqual(0, models.Job.objects.all().count())
+
+
+    def test_normal_delete_queryset(self):
+        self._create_job(hosts=[1])
+        self._create_job(hosts=[2])
+
+        self.assertEqual(2, models.Job.objects.all().count())
+
+        models.Job.objects.all().delete()
+        self.assertEqual(0, models.Job.objects.all().count())
+
+
 if __name__ == '__main__':
     unittest.main()
_______________________________________________
Autotest mailing list
[email protected]
http://test.kernel.org/cgi-bin/mailman/listinfo/autotest

Reply via email to