Fix leap year unittest failures. (#4953)

Author: myersCody

koku/koku/test_pg_partition.py

@@ -2,9 +2,11 @@
 # Copyright 2021 Red Hat Inc.
 # SPDX-License-Identifier: Apache-2.0
 #
+import calendar
 import datetime
 import json
 import uuid
+from datetime import timedelta
 
 from django.db import connection as conn
 from django_tenants.utils import schema_context
@@ -22,6 +24,15 @@ def _execute(sql, params=None):
     return cur
 
 
+def handle_leap_year_overflow(date, years):
+    try:
+        new_date = date.replace(year=(date.year + years))
+    except ValueError:
+        # If an error occurs, set the day to the last day of the month
+        new_date = (date + timedelta(days=31)).replace(year=(date.year + years), day=1)
+    return new_date
+
+
 def _get_table_partition_info(schema, table):
     sql = """
 select c.relkind,
@@ -909,7 +920,7 @@ def test_repartition_table(self):
         """
         with schema_context(self.schema_name):
             aws_lids = AWSCostEntryLineItemDailySummary.objects.order_by("-usage_start")[0]
-            year10_usage_start = aws_lids.usage_start.replace(year=(aws_lids.usage_start.year + 10))
+            year10_usage_start = handle_leap_year_overflow(aws_lids.usage_start, 10)
             aws_lids.usage_start = year10_usage_start
             aws_lids.save()
             year10_count = 0
@@ -952,10 +963,11 @@ def test_repartition_all_tables(self):
         """
         with schema_context(self.schema_name):
             aws_lids = AWSCostEntryLineItemDailySummary.objects.order_by("-usage_start")[0]
+            aws_lids.usage_start = handle_leap_year_overflow(aws_lids.usage_start, 11)
             aws_lids.usage_start = aws_lids.usage_start.replace(year=(aws_lids.usage_start.year + 11))
             aws_lids.save()
             ocp_lids = OCPUsageLineItemDailySummary.objects.order_by("-usage_start")[0]
-            ocp_lids.usage_start = ocp_lids.usage_start.replace(year=(aws_lids.usage_start.year + 11))
+            ocp_lids.usage_start = handle_leap_year_overflow(aws_lids.usage_start, 11)
             ocp_lids.save()
             with conn.cursor() as cur:
                 cur.execute(
@@ -1019,7 +1031,14 @@ def test_repartition_all_tables(self):
                 if new_ocp_lids.usage_start.day < 28
                 else new_ocp_lids.usage_start.day - 1
             )
-            new_ocp_lids.usage_start = ocp_lids.usage_start.replace(day=new_day)
+            if (
+                ocp_lids.usage_start.month == 2
+                and ocp_lids.usage_start.day == 28
+                and calendar.isleap(ocp_lids.usage_start.year)
+            ):
+                new_ocp_lids.usage_start = ocp_lids.usage_start.replace(day=29)
+            else:
+                new_ocp_lids.usage_start = ocp_lids.usage_start.replace(day=new_day)
             new_ocp_lids.save()
 
             with conn.cursor() as cur: