Merge pull request #231 from American-Institutes-for-Research/HEA-914/Add-as_of_date-filter-to-the-LivelihoodZoneBaselineViewSet

Add as_of_date filer to LivelihoodZoneBaselineFilterSet and correspon…

Author: rhunwicks

apps/baseline/models.py

@@ -2,6 +2,7 @@
 Models for managing HEA Baseline Surveys
 """
 
+import datetime
 import numbers
 
 from django.conf import settings
@@ -131,7 +132,30 @@ class ExtraMeta:
         identifier = ["code"]
 
 
-class LivelihoodZoneBaselineManager(common_models.IdentifierManager):
+class LivelihoodZoneBaselineQuerySet(models.QuerySet):
+    """
+    QuerySet for LivelihoodZoneBaseline that provides temporal filtering methods.
+    """
+
+    def filter_current(self, as_of_date=None):
+        """
+        Return a queryset filtered to the baselines that are valid as of the date specified.
+        """
+        if not as_of_date:
+            as_of_date = datetime.date.today()
+        return self.filter(
+            (models.Q(valid_from_date__lte=as_of_date) | models.Q(valid_from_date__isnull=True))
+            & (models.Q(valid_to_date__gte=as_of_date) | models.Q(valid_to_date__isnull=True))
+        )
+
+    def current_all(self, as_of_date=None):
+        """
+        Return all the baselines that are valid as of the date specified.
+        """
+        return self.filter_current(as_of_date).all()
+
+
+class LivelihoodZoneBaselineManager(common_models.IdentifierManager.from_queryset(LivelihoodZoneBaselineQuerySet)):
     def get_by_natural_key(self, code: str, reference_year_end_date: str):
         return self.get(livelihood_zone__code=code, reference_year_end_date=reference_year_end_date)
 

apps/baseline/tests/factories.py

@@ -119,8 +119,9 @@ class Meta:
     bss_language = factory.Iterator(["en", "pt", "es", "ar", "fr"])
     reference_year_start_date = factory.LazyAttribute(lambda o: o.reference_year_end_date - relativedelta(years=1))
     reference_year_end_date = factory.Sequence(lambda n: datetime.date(1900, 1, 1) + datetime.timedelta(days=n + 10))
-    valid_from_date = factory.Sequence(lambda n: datetime.date(1900, 1, 1) + datetime.timedelta(days=n))
-    valid_to_date = factory.Sequence(lambda n: datetime.date(1900, 1, 1) + datetime.timedelta(days=n + 10))
+    # Default to None so factory-created baselines are always valid regardless of as_of_date filter
+    valid_from_date = None
+    valid_to_date = None
     population_source = factory.Sequence(lambda n: f"population_source {n}")
     population_estimate = fuzzy.FuzzyInteger(500, 1000000)
     currency = factory.SubFactory(CurrencyFactory)

apps/baseline/tests/test_viewsets.py

@@ -595,6 +595,88 @@ def test_filter_by_wealth_characteristic(self):
         self.assertEqual(response.status_code, 200)
         self.assertEqual(len(response.json()), 1)
 
+    def test_as_of_date_filter_returns_valid_baselines(self):
+        """
+        Test that the as_of_date filter returns only baselines valid as of the specified date.
+        """
+        today = datetime.date.today()
+
+        # Baseline that expired in the past
+        expired_baseline = LivelihoodZoneBaselineFactory(
+            valid_from_date=today - datetime.timedelta(days=365),
+            valid_to_date=today - datetime.timedelta(days=30),
+        )
+
+        # Baseline currently valid
+        current_baseline = LivelihoodZoneBaselineFactory(
+            valid_from_date=today - datetime.timedelta(days=30),
+            valid_to_date=today + datetime.timedelta(days=365),
+        )
+
+        # Baseline that starts in the future (not yet valid)
+        future_baseline = LivelihoodZoneBaselineFactory(
+            valid_from_date=today + datetime.timedelta(days=30),
+            valid_to_date=today + datetime.timedelta(days=365),
+        )
+
+        # Test default behavior (no as_of_date param) - should return all baselines
+        response = self.client.get(self.url)
+        self.assertEqual(response.status_code, 200)
+        baseline_ids = [b["id"] for b in response.json()]
+        self.assertIn(current_baseline.id, baseline_ids)
+        self.assertIn(expired_baseline.id, baseline_ids)
+        self.assertIn(future_baseline.id, baseline_ids)
+
+        # Test with explicit as_of_date=today - should filter to current baselines
+        response = self.client.get(self.url, {"as_of_date": today.isoformat()})
+        self.assertEqual(response.status_code, 200)
+        baseline_ids = [b["id"] for b in response.json()]
+        self.assertIn(current_baseline.id, baseline_ids)
+        self.assertNotIn(expired_baseline.id, baseline_ids)
+        self.assertNotIn(future_baseline.id, baseline_ids)
+
+        # Test with a past date when expired_baseline was still valid
+        # current_baseline hadn't started yet, so it should be excluded
+        past_date = today - datetime.timedelta(days=180)
+        response = self.client.get(self.url, {"as_of_date": past_date.isoformat()})
+        self.assertEqual(response.status_code, 200)
+        baseline_ids = [b["id"] for b in response.json()]
+        self.assertIn(expired_baseline.id, baseline_ids)
+        self.assertNotIn(current_baseline.id, baseline_ids)  # hadn't started yet
+        self.assertNotIn(future_baseline.id, baseline_ids)
+
+        # Test with a future date when future_baseline will be valid
+        future_date = today + datetime.timedelta(days=60)
+        response = self.client.get(self.url, {"as_of_date": future_date.isoformat()})
+        self.assertEqual(response.status_code, 200)
+        baseline_ids = [b["id"] for b in response.json()]
+        self.assertNotIn(expired_baseline.id, baseline_ids)
+        self.assertIn(current_baseline.id, baseline_ids)
+        self.assertIn(future_baseline.id, baseline_ids)
+
+        # Baseline with null valid_from_date (valid from the beginning of time)
+        baseline_no_from = LivelihoodZoneBaselineFactory(
+            valid_from_date=None,
+            valid_to_date=today + datetime.timedelta(days=365),
+        )
+        # Baseline with null valid_to_date (valid indefinitely)
+        baseline_no_to = LivelihoodZoneBaselineFactory(
+            valid_from_date=today - datetime.timedelta(days=365),
+            valid_to_date=None,
+        )
+        # Baseline with both dates null (always valid)
+        baseline_no_dates = LivelihoodZoneBaselineFactory(
+            valid_from_date=None,
+            valid_to_date=None,
+        )
+        # Test default behavior - all three should be returned
+        response = self.client.get(self.url)
+        self.assertEqual(response.status_code, 200)
+        baseline_ids = [b["id"] for b in response.json()]
+        self.assertIn(baseline_no_from.id, baseline_ids)
+        self.assertIn(baseline_no_to.id, baseline_ids)
+        self.assertIn(baseline_no_dates.id, baseline_ids)
+
 
 class LivelihoodZoneBaselineFacetedSearchViewTestCase(APITestCase):
     def setUp(self):
@@ -641,7 +723,8 @@ def test_search_with_product(self):
         # Apply the filters to the baseline
         baseline_url = reverse("livelihoodzonebaseline-list")
         response = self.client.get(
-            baseline_url, {search_data["products"][0]["filter"]: search_data["products"][0]["value"]}
+            baseline_url,
+            {search_data["products"][0]["filter"]: search_data["products"][0]["value"]},
         )
         self.assertEqual(response.status_code, 200)
         self.assertEqual(len(json.loads(response.content)), 2)
@@ -650,7 +733,10 @@ def test_search_with_product(self):
         self.assertTrue(any(d["name"] == self.baseline3.name for d in data))
         self.assertFalse(any(d["name"] == self.baseline2.name for d in data))
 
-        response = self.client.get(baseline_url, {search_data["items"][0]["filter"]: search_data["items"][0]["value"]})
+        response = self.client.get(
+            baseline_url,
+            {search_data["items"][0]["filter"]: search_data["items"][0]["value"]},
+        )
         self.assertEqual(response.status_code, 200)
         self.assertEqual(len(json.loads(response.content)), 1)
         data = json.loads(response.content)

apps/baseline/viewsets.py

@@ -13,7 +13,7 @@
 from rest_framework.views import APIView
 
 from common.fields import translation_fields
-from common.filters import MultiFieldFilter, UpperCaseFilter
+from common.filters import DefaultingDateFilter, MultiFieldFilter, UpperCaseFilter
 from common.viewsets import AggregatingViewSet, BaseModelViewSet
 from metadata.models import WealthGroupCategory
 
@@ -205,6 +205,10 @@ class Meta:
     wealth_characteristic = CharFilter(
         method="filter_by_wealth_characteristic", label="Filter by Wealth Characteristic"
     )
+    as_of_date = DefaultingDateFilter(
+        label="As of Date",
+        help_text="Filter baselines valid as of this date (YYYY-MM-DD format or special values like 'today').",
+    )
 
     def filter_by_product(self, queryset, name, value):
         """

apps/common/filters.py

@@ -8,14 +8,16 @@
 from django.core.validators import EMPTY_VALUES
 from django.db.models import F, Func, Q
 from django.forms import TextInput
-from django.forms.fields import ChoiceField, Field, MultipleChoiceField
+from django.forms.fields import ChoiceField, DateField, Field, MultipleChoiceField
 from django.forms.models import ModelMultipleChoiceField
 from django.utils.datastructures import MultiValueDict
 from django.utils.encoding import force_str
 from django_filters import ModelMultipleChoiceFilter, MultipleChoiceFilter
-from django_filters.filters import BooleanFilter, CharFilter, ChoiceFilter
+from django_filters.filters import BooleanFilter, CharFilter, ChoiceFilter, DateFilter
 from rest_framework.filters import OrderingFilter
 
+from .utils import DEFAULT_DATES
+
 logger = logging.getLogger(__name__)
 
 
@@ -238,3 +240,33 @@ def filter(self, qs, value):
         method = qs.exclude if exclude else qs.filter
 
         return method(**{self.field_name: ""})
+
+
+class DefaultingDateField(DateField):
+    """
+    A date field that accepts defaults like "last_month", "today"
+    """
+
+    def to_python(self, value):
+        if value in DEFAULT_DATES:
+            value = DEFAULT_DATES[value]()
+        return super().to_python(value)
+
+
+class DefaultingDateFilter(DateFilter):
+    """
+    A date filter that accepts defaults like "today"
+    """
+
+    field_class = DefaultingDateField
+
+    def filter(self, qs, value):
+        if value:
+            if self.lookup_expr in ["lte", "gte"]:
+                # period_date filter for start and end date with either gte or lte expression will fall here.
+                query = Q()
+                query = Q(**{self.field_name + "__" + self.lookup_expr: value})
+                return qs.filter(query)
+            else:
+                return qs.current_all(value)
+        return qs

apps/common/utils.py

@@ -9,6 +9,7 @@
 from pathlib import Path
 
 import pandas as pd
+from dateutil.relativedelta import relativedelta
 from django.apps import apps
 from django.db.migrations.operations.base import Operation
 from django.forms.models import modelform_factory
@@ -18,6 +19,24 @@
 logger = logging.getLogger(__name__)
 
 
+DEFAULT_DATES = {
+    "ten_years_ago": lambda: datetime.date.today() + relativedelta(years=-10, day=1),
+    "five_years_ago": lambda: datetime.date.today() + relativedelta(years=-5, day=1),
+    "three_years_ago": lambda: datetime.date.today() + relativedelta(years=-3, day=1),
+    "two_years_ago": lambda: datetime.date.today() + relativedelta(years=-2, day=1),
+    "one_year_ago": lambda: datetime.date.today() + relativedelta(years=-1, day=1),
+    "last_month_start": lambda: datetime.date.today() + relativedelta(months=-1, day=1),
+    "last_month": lambda: datetime.date.today().replace(day=1) - datetime.timedelta(days=1),
+    "second_last_month_start": lambda: datetime.date.today() + relativedelta(months=-2, day=1),
+    "second_last_month": lambda: datetime.date.today() + relativedelta(months=-2, day=31),
+    "third_last_month": lambda: datetime.date.today() + relativedelta(months=-3, day=31),
+    "fourth_last_month": lambda: datetime.date.today() + relativedelta(months=-4, day=31),
+    "sixth_last_month": lambda: datetime.date.today() + relativedelta(months=-6, day=31),
+    "today": lambda: datetime.date.today(),
+    "tomorrow": lambda: datetime.date.today() + relativedelta(days=1),
+}
+
+
 class UnicodeCsvReader(object):
     # TODO: Should check if it works without encoding and if this class needed on Python 3 at all.
     def __init__(self, f, encoding="utf-8", **kwargs):