Custom QuerySet

custom_queryset_test.py

import unittest
from datetime import datetime
from mongoengine import *
from mongoengine.queryset import QuerySet


class BugFixTest(unittest.TestCase):

    def setUp(self):
        conn = connect(db='mongoenginetest')

    def test_custom_query_set(self):

        class CustomQuerySet(QuerySet):

            def get_recent_by_date(self, date):
                """
                gets recent activity from the start date given, to the
                current time
                """
                return self.filter(created_time__gte=date)

            def get_recent_by_user(self, user):
                """
                gets activity from the given user
                """
                return self.filter(performed_by=user)

            def get_activity(self, start_date, end_date):
                """
                gets all activity from the start date to the end date
                """
                return self.filter(created_time__lt=end_date,
                                   created_time__gte=start_date)

        class Doc(Document):
            performed_by = StringField()
            created_time = DateTimeField()

            meta = {'queryset_class': CustomQuerySet}

        Doc.drop_collection()
        Doc(performed_by="Me", created_time=datetime.now()).save()
        Doc(performed_by="You", created_time=datetime(2010, 01, 01)).save()

        self.assertEqual(1, Doc.objects.get_recent_by_date(
                                datetime(2011, 01, 01)).count())
        self.assertEqual(1, Doc.objects.get_recent_by_user("You").count())
        self.assertEqual(1, Doc.objects.get_activity(
                                datetime(2009, 01, 01),
                                datetime(2011, 01, 01)).count())