[AIRFLOW-6165] Housekeep utils.dates.date_range & add tests (#6720)

airflow/utils/dates.py

@@ -59,6 +59,8 @@ def date_range(start_date, end_date=None, num=None, delta=None):
         number of entries you want in the range. This number can be negative,
         output will always be sorted regardless
     :type num: int
+    :param delta: step length. It can be datetime.timedelta or cron expression as string
+    :type delta: datetime.timedelta or str
     """
     if not delta:
         return []
@@ -71,12 +73,17 @@ def date_range(start_date, end_date=None, num=None, delta=None):
 
     delta_iscron = False
     tz = start_date.tzinfo
+
     if isinstance(delta, str):
         delta_iscron = True
-        start_date = timezone.make_naive(start_date, tz)
+        if timezone.is_localized(start_date):
+            start_date = timezone.make_naive(start_date, tz)
         cron = croniter(delta, start_date)
     elif isinstance(delta, timedelta):
         delta = abs(delta)
+    else:
+        raise Exception("Wait. delta must be either datetime.timedelta or cron expression as str")
+
     dates = []
     if end_date:
         if timezone.is_naive(start_date) and not timezone.is_naive(end_date):

tests/utils/test_dates.py

@@ -52,3 +52,46 @@ class TestDates(unittest.TestCase):
             timezone.datetime(2017, 11, 5, 16, 18, 30, 989729),
             dates.parse_execution_date(execution_date_str_w_ms))
         self.assertRaises(ValueError, dates.parse_execution_date, bad_execution_date_str)
+
+
+class TestUtilsDatesDateRange(unittest.TestCase):
+
+    def test_no_delta(self):
+        self.assertEqual(dates.date_range(datetime(2016, 1, 1), datetime(2016, 1, 3)),
+                         [])
+
+    def test_end_date_before_start_date(self):
+        with self.assertRaisesRegex(Exception, "Wait. start_date needs to be before end_date"):
+            dates.date_range(datetime(2016, 2, 1),
+                             datetime(2016, 1, 1),
+                             delta=timedelta(seconds=1))
+
+    def test_both_end_date_and_num_given(self):
+        with self.assertRaisesRegex(Exception, "Wait. Either specify end_date OR num"):
+            dates.date_range(datetime(2016, 1, 1),
+                             datetime(2016, 1, 3),
+                             num=2,
+                             delta=timedelta(seconds=1))
+
+    def test_invalid_delta(self):
+        exception_msg = "Wait. delta must be either datetime.timedelta or cron expression as str"
+        with self.assertRaisesRegex(Exception, exception_msg):
+            dates.date_range(datetime(2016, 1, 1),
+                             datetime(2016, 1, 3),
+                             delta=1)
+
+    def test_positive_num_given(self):
+        for num in range(1, 10):
+            result = dates.date_range(datetime(2016, 1, 1), num=num, delta=timedelta(1))
+            self.assertEqual(len(result), num)
+
+            for i in range(num):
+                self.assertTrue(timezone.is_localized(result[i]))
+
+    def test_negative_num_given(self):
+        for num in range(-1, -5, -10):
+            result = dates.date_range(datetime(2016, 1, 1), num=num, delta=timedelta(1))
+            self.assertEqual(len(result), -num)
+
+            for i in range(num):
+                self.assertTrue(timezone.is_localized(result[i]))