Tests for PassThroughManager

django_pandas/tests/models.py

@@ -1,6 +1,6 @@
 from django.db import models
 from django.utils.encoding import python_2_unicode_compatible
-from django_pandas.managers import DataFrameManager
+from django_pandas.managers import DataFrameManager, PassThroughManager
 
 
 @python_2_unicode_compatible
@@ -158,3 +158,71 @@ class Portfolio(models.Model):
 
     def __str__(self):
         return self.name
+
+
+class DudeQuerySet(models.query.QuerySet):
+    def abiding(self):
+        return self.filter(abides=True)
+
+    def rug_positive(self):
+        return self.filter(has_rug=True)
+
+    def rug_negative(self):
+        return self.filter(has_rug=False)
+
+    def by_name(self, name):
+        return self.filter(name__iexact=name)
+
+
+class AbidingManager(PassThroughManager):
+    def get_queryset(self):
+        return DudeQuerySet(self.model).abiding()
+
+    get_query_set = get_queryset
+
+    def get_stats(self):
+        return {
+            "abiding_count": self.count(),
+            "rug_count": self.rug_positive().count(),
+        }
+
+
+class Dude(models.Model):
+    abides = models.BooleanField(default=True)
+    name = models.CharField(max_length=20)
+    has_rug = models.BooleanField(default=False)
+
+    objects = PassThroughManager(DudeQuerySet)
+    abiders = AbidingManager()
+
+
+class Car(models.Model):
+    name = models.CharField(max_length=20)
+    owner = models.ForeignKey(Dude, related_name='cars_owned')
+
+    objects = PassThroughManager(DudeQuerySet)
+
+
+class SpotManager(PassThroughManager):
+    def get_queryset(self):
+        return super(SpotManager, self).get_queryset().filter(secret=False)
+
+    get_query_set = get_queryset
+
+
+class SpotQuerySet(models.query.QuerySet):
+    def closed(self):
+        return self.filter(closed=True)
+
+    def secured(self):
+        return self.filter(secure=True)
+
+
+class Spot(models.Model):
+    name = models.CharField(max_length=20)
+    secure = models.BooleanField(default=True)
+    closed = models.BooleanField(default=False)
+    secret = models.BooleanField(default=False)
+    owner = models.ForeignKey(Dude, related_name='spots_owned')
+
+    objects = SpotManager.for_queryset_class(SpotQuerySet)()

django_pandas/tests/test_manager.py

@@ -1,9 +1,11 @@
 from django.test import TestCase
 import pandas as pd
 import numpy as np
+import pickle
+import django
 from .models import (
     DataFrame, WideTimeSeries,
-    LongTimeSeries, PivotData, MyModelChoice
+    LongTimeSeries, PivotData, MyModelChoice, Dude, Car, Spot
 )
 import pandas.util.testing as tm
 
@@ -183,3 +185,114 @@ def test_pivot(self):
         self.assertEqual(pt.index.names, rows)
         self.assertEqual(pt.columns.names, cols)
 
+
+if django.VERSION < (1, 9):
+
+    class PassThroughManagerTests(TestCase):
+
+        def setUp(self):
+            Dude.objects.create(name='The Dude', abides=True, has_rug=False)
+            Dude.objects.create(name='His Dudeness', abides=False, has_rug=True)
+            Dude.objects.create(name='Duder', abides=False, has_rug=False)
+            Dude.objects.create(name='El Duderino', abides=True, has_rug=True)
+
+
+        def test_chaining(self):
+            self.assertEqual(Dude.objects.by_name('Duder').count(), 1)
+            self.assertEqual(Dude.objects.all().by_name('Duder').count(), 1)
+            self.assertEqual(Dude.abiders.rug_positive().count(), 1)
+            self.assertEqual(Dude.abiders.all().rug_positive().count(), 1)
+
+
+        def test_manager_only_methods(self):
+            stats = Dude.abiders.get_stats()
+            self.assertEqual(stats['rug_count'], 1)
+            with self.assertRaises(AttributeError):
+                Dude.abiders.all().get_stats()
+
+
+        def test_queryset_pickling(self):
+            qs = Dude.objects.all()
+            saltyqs = pickle.dumps(qs)
+            unqs = pickle.loads(saltyqs)
+            self.assertEqual(unqs.by_name('The Dude').count(), 1)
+
+
+        def test_queryset_not_available_on_related_manager(self):
+            dude = Dude.objects.by_name('Duder').get()
+            Car.objects.create(name='Ford', owner=dude)
+            self.assertFalse(hasattr(dude.cars_owned, 'by_name'))
+
+
+        def test_using_dir(self):
+            # make sure introspecing via dir() doesn't actually cause queries,
+            # just as a sanity check.
+            with self.assertNumQueries(0):
+                querysets_to_dir = (
+                    Dude.objects,
+                    Dude.objects.by_name('Duder'),
+                    Dude.objects.all().by_name('Duder'),
+                    Dude.abiders,
+                    Dude.abiders.rug_positive(),
+                    Dude.abiders.all().rug_positive()
+                )
+                for qs in querysets_to_dir:
+                    self.assertTrue('by_name' in dir(qs))
+                    self.assertTrue('abiding' in dir(qs))
+                    self.assertTrue('rug_positive' in dir(qs))
+                    self.assertTrue('rug_negative' in dir(qs))
+                    # some standard qs methods
+                    self.assertTrue('count' in dir(qs))
+                    self.assertTrue('order_by' in dir(qs))
+                    self.assertTrue('select_related' in dir(qs))
+                    # make sure it's been de-duplicated
+                    self.assertEqual(1, dir(qs).count('distinct'))
+
+                # manager only method.
+                self.assertTrue('get_stats' in dir(Dude.abiders))
+                # manager only method shouldn't appear on the non AbidingManager
+                self.assertFalse('get_stats' in dir(Dude.objects))
+                # standard manager methods
+                self.assertTrue('get_query_set' in dir(Dude.abiders))
+                self.assertTrue('contribute_to_class' in dir(Dude.abiders))
+
+
+
+    class CreatePassThroughManagerTests(TestCase):
+
+        def setUp(self):
+            self.dude = Dude.objects.create(name='El Duderino')
+            self.other_dude = Dude.objects.create(name='Das Dude')
+
+        def test_reverse_manager(self):
+            Spot.objects.create(
+                name='The Crib', owner=self.dude, closed=True, secure=True,
+                secret=False)
+            self.assertEqual(self.dude.spots_owned.closed().count(), 1)
+            Spot.objects.create(
+                name='The Crux', owner=self.other_dude, closed=True, secure=True,
+                secret=False
+            )
+            self.assertEqual(self.dude.spots_owned.closed().all().count(), 1)
+            self.assertEqual(self.dude.spots_owned.closed().count(), 1)
+
+        def test_related_queryset_pickling(self):
+            Spot.objects.create(
+                name='The Crib', owner=self.dude, closed=True, secure=True,
+                secret=False)
+            qs = self.dude.spots_owned.closed()
+            pickled_qs = pickle.dumps(qs)
+            unpickled_qs = pickle.loads(pickled_qs)
+            self.assertEqual(unpickled_qs.secured().count(), 1)
+
+        def test_related_queryset_superclass_method(self):
+            Spot.objects.create(
+                name='The Crib', owner=self.dude, closed=True, secure=True,
+                secret=False)
+            Spot.objects.create(
+                name='The Secret Crib', owner=self.dude, closed=False, secure=True,
+                secret=True)
+            self.assertEqual(self.dude.spots_owned.count(), 1)
+
+        def test_related_manager_create(self):
+            self.dude.spots_owned.create(name='The Crib', closed=True, secure=True)