[DevTest Labs] az lab vm create: Fix --expiration-date always failing with a datetime comparison error

src/azure-cli/azure/cli/command_modules/lab/tests/latest/test_lab_validators.py

@@ -3,10 +3,14 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+import datetime
 import unittest
 from azure.mgmt.core.tools import is_valid_resource_id
+from azure.cli.core.aaz import AAZArgumentsSchema, AAZDateTimeArg
+from azure.cli.core.aaz._command_ctx import AAZCommandCtx
 from azure.cli.core.azclierror import ArgumentUsageError
-from azure.cli.command_modules.lab.validators import _update_artifacts
+from azure.cli.core.mock import DummyCli
+from azure.cli.command_modules.lab.validators import _update_artifacts, _validate_expiration_date
 
 
 class ValidatorsCommandTest(unittest.TestCase):
@@ -65,3 +69,31 @@ def test_update_artifacts(self):
         del invalid_artifact['artifact_id']
         with self.assertRaises(ArgumentUsageError):
             _update_artifacts([invalid_artifact], self.lab_resource_id)
+
+
+class ExpirationDateValidatorTest(unittest.TestCase):
+    @staticmethod
+    def _build_args(expiration_date=None):
+        """ Builds the args `az lab vm create` hands to the validator. """
+        schema = AAZArgumentsSchema()
+        schema.expiration_date = AAZDateTimeArg(options=['--expiration-date'])
+        command_args = {} if expiration_date is None else {'expiration_date': expiration_date}
+        ctx = AAZCommandCtx(cli_ctx=DummyCli(), schema=schema, command_args=command_args)
+        ctx.format_args()
+        return ctx.args
+
+    def test_expiration_date_in_future(self):
+        future = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=2)
+        # No offset, a Z suffix and an explicit offset all reach the validator as UTC timestamps.
+        for expiration_date in (future.strftime('%Y-%m-%d %H:%M:%S'),
+                                future.strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
+                                future.strftime('%Y-%m-%d %H:%M:%S+00:00')):
+            _validate_expiration_date(self._build_args(expiration_date))
+
+    def test_expiration_date_in_past(self):
+        past = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=2)
+        with self.assertRaises(ArgumentUsageError):
+            _validate_expiration_date(self._build_args(past.strftime('%Y-%m-%dT%H:%M:%S.%fZ')))
+
+    def test_expiration_date_not_provided(self):
+        _validate_expiration_date(self._build_args())

src/azure-cli/azure/cli/command_modules/lab/validators.py

@@ -115,7 +115,10 @@ def _validate_expiration_date(args):
     if has_value(args.expiration_date):
         import datetime
         import dateutil.parser
-        if datetime.datetime.utcnow() >= dateutil.parser.parse(args.expiration_date.to_serialized_data()):
+        # AAZDateTimeArg normalizes --expiration-date into an offset-aware UTC timestamp, so "now"
+        # has to be offset-aware too. datetime.utcnow() is naive and the comparison raises TypeError.
+        if datetime.datetime.now(datetime.timezone.utc) >= \
+                dateutil.parser.parse(args.expiration_date.to_serialized_data()):
             raise ArgumentUsageError(
                 "Expiration date '{}' must be in future.".format(args.expiration_date.to_serialized_data()))