Add QuerySet methods to check data expiration

wifiwithme/apps/contribmap/models.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 from __future__ import unicode_literals
+from datetime import timedelta
 
 from django.core.urlresolvers import reverse
 from django.db import models
@@ -10,6 +11,24 @@ from .fields import CommaSeparatedCharField
 from .utils import add_one_year, ANGLES, merge_intervals
 
 
+class ContribQuerySet(models.query.QuerySet):
+    def expired(self):
+        return self.filter(expiration_date__lte=timezone.now())
+
+    def expired_in_days(self, ndays):
+        return self.filter(
+            expiration_date__lte=timezone.now() + timedelta(days=ndays))
+
+    def expires_in_days(self, ndays):
+        """ Returns only the data expiring in exactly that number of days
+
+        Think about it as an anniversary. This function ignores
+        minutes/seconds/hours and check only days.
+        """
+        return self.filter(
+            expiration_date__date=timezone.now() + timedelta(days=ndays))
+
+
 class Contrib(models.Model):
     CONTRIB_CONNECT = 'connect'
     CONTRIB_SHARE = 'share'
@@ -110,6 +129,8 @@ class Contrib(models.Model):
 
     PUBLIC_FIELDS = set(PRIVACY_MAP.keys())
 
+    objects = ContribQuerySet.as_manager()
+
     def __str__(self):
         return '#{} {}'.format(self.pk, self.name)
 

wifiwithme/apps/contribmap/tests.py

@@ -89,6 +89,50 @@ class TestContribPrivacy(TestCase):
         self.assertEqual(c.get_public_field('name'), None)
 
 
+class TestContribQuerySet(TestCase):
+    def test_expired(self):
+        with freeze_time('12-11-2100', tz_offset=0):
+            Contrib.objects.create(
+                name='foo', orientations=['S'],
+                contrib_type=Contrib.CONTRIB_CONNECT,
+                latitude=0.5, longitude=0.5)
+        # one year and one month later
+        with freeze_time('12-12-2101', tz_offset=0):
+            Contrib.objects.create(
+                name='bar', orientations=['S'],
+                contrib_type=Contrib.CONTRIB_CONNECT,
+                latitude=0.5, longitude=0.5)
+
+            expired = Contrib.objects.expired()
+            self.assertEqual(expired.count(), 1)
+            self.assertEqual(expired.first().name, 'foo')
+
+    def test_expired_in_days(self):
+        Contrib.objects.create(
+            name='foo', orientations=['S'],
+            contrib_type=Contrib.CONTRIB_CONNECT,
+            latitude=0.5, longitude=0.5)
+
+        self.assertEqual(Contrib.objects.expired_in_days(0).count(), 0)
+        self.assertEqual(Contrib.objects.expired_in_days(366).count(), 1)
+
+    def test_expires_in_days(self):
+        with freeze_time('12-11-2101 12:00', tz_offset=0):
+            Contrib.objects.create(
+                name='foo', orientations=['S'],
+                contrib_type=Contrib.CONTRIB_CONNECT,
+                latitude=0.5, longitude=0.5)
+            self.assertEqual(Contrib.objects.expires_in_days(364).count(), 0)
+            self.assertEqual(Contrib.objects.expires_in_days(365).count(), 1)
+            self.assertEqual(Contrib.objects.expires_in_days(366).count(), 0)
+
+        # One year, one hour and two minutes later
+        # (check that minutes/hours are ignored)
+        with freeze_time('12-11-2101 13:02', tz_offset=0):
+            self.assertEqual(Contrib.objects.expires_in_days(365).count(), 1)
+
+
+
 class TestViews(APITestCase):
     def mk_contrib_post_data(self, *args, **kwargs):
         post_data = {