Allow "now" as a valid datetime argument- #2

dftimewolf/lib/validators/datetime_validator.py

@@ -9,11 +9,11 @@
 from dftimewolf.lib.validators import manager as validators_manager
 
 
-
 class DatetimeValidator(args_validator.AbstractValidator):
   """Validates a date and time string.
 
-  Accepts dates in ISO2601 format only.
+  Accepts dates in ISO8601 format, and the special value ``now`` for the
+  current UTC time.
   """
 
   NAME = 'datetime'
@@ -36,6 +36,24 @@ def _EnsureTimezone(
       parsed_datetime = parsed_datetime.replace(tzinfo=datetime.timezone.utc)
     return parsed_datetime
 
+  def _ParseDatetime(self, value: str) -> datetime.datetime:
+    """Parses a datetime string, including the special value ``now``.
+
+    Args:
+      value: The datetime string to parse.
+
+    Returns:
+      A timezone-aware datetime object.
+
+    Raises:
+      parser.ParserError: If the value is not a valid datetime.
+      ValueError: If the value is not a valid datetime.
+    """
+    if value.lower() == 'now':
+      return datetime.datetime.now(datetime.timezone.utc)
+    parsed_datetime = parser.isoparse(value)
+    return self._EnsureTimezone(parsed_datetime)
+
   def Validate(self, argument_value: Any,
                recipe_argument: resources.RecipeArgument) -> datetime.datetime:
     """Validate that operand is a valid date and time string.
@@ -62,8 +80,7 @@ def Validate(self, argument_value: Any,
     validation_parameters = recipe_argument.validation_params
 
     try:
-      parsed_datetime = parser.isoparse(argument_value)
-      parsed_datetime = self._EnsureTimezone(parsed_datetime)
+      parsed_datetime = self._ParseDatetime(argument_value)
     except (parser.ParserError, ValueError) as exception:
       raise errors.RecipeArgsValidationFailure(
           recipe_argument.switch,
@@ -118,13 +135,11 @@ def _ValidateOrder(self,
     first_datetime: datetime.datetime
     second_datetime: datetime.datetime
     if isinstance(first, str):
-      first_datetime = parser.isoparse(first)
-      first_datetime = self._EnsureTimezone(first_datetime)
+      first_datetime = self._ParseDatetime(first)
     else:
       first_datetime = first
     if isinstance(second, str):
-      second_datetime = parser.isoparse(second)
-      second_datetime = self._EnsureTimezone(second_datetime)
+      second_datetime = self._ParseDatetime(second)
     else:
       second_datetime = second
 

tests/lib/validators/datetime_validator.py

@@ -3,6 +3,7 @@
 
 import datetime
 import unittest
+from unittest import mock
 
 from absl.testing import parameterized
 
@@ -36,6 +37,20 @@ def testValidateSuccess(self):
     val = self.validator.Validate(date_string, self.recipe_argument)
     self.assertEqual(val, date_value)
 
+  def testValidateNow(self):
+    """Tests that ``now`` resolves to the current UTC time."""
+    frozen_now = datetime.datetime(
+        2024, 6, 15, 12, 30, 45, tzinfo=datetime.timezone.utc)
+    with mock.patch.object(datetime_validator, 'datetime') as mock_datetime:
+      mock_datetime.datetime.now.return_value = frozen_now
+      mock_datetime.timezone = datetime.timezone
+      for now_value in ('now', 'NOW', 'Now'):
+        with self.subTest(now_value=now_value):
+          val = self.validator.Validate(now_value, self.recipe_argument)
+          self.assertEqual(val, frozen_now)
+          mock_datetime.datetime.now.assert_called_with(
+              datetime.timezone.utc)
+
   def testValidateSuccessWithOrder(self):
     """Tests validation success with order parameters."""
     first_string = '2023-01-01 00:00:00'