From dac23d77a349c8fe025da871b0bda6222879dd9e Mon Sep 17 00:00:00 2001
From: Zack Cerza <zack@redhat.com>
Date: Wed, 31 Jul 2024 14:11:52 -0600
Subject: [PATCH] util.time: Add parse_timestamp

And move the format string to the time module.

Signed-off-by: Zack Cerza <zack@redhat.com>
---
 teuthology/suite/run.py            |  3 ++-
 teuthology/suite/test/test_run_.py |  3 ++-
 teuthology/util/test/test_time.py  | 26 +++++++++++++++++++++++++-
 teuthology/util/time.py            | 20 +++++++++++++++++++-
 4 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/teuthology/suite/run.py b/teuthology/suite/run.py
index 8fa45b5b..600f1568 100644
--- a/teuthology/suite/run.py
+++ b/teuthology/suite/run.py
@@ -24,6 +24,7 @@ from teuthology.suite import util
 from teuthology.suite.merge import config_merge
 from teuthology.suite.build_matrix import build_matrix
 from teuthology.suite.placeholder import substitute_placeholders, dict_templ
+from teuthology.util.time import TIMESTAMP_FMT
 
 log = logging.getLogger(__name__)
 
@@ -43,7 +44,7 @@ class Run(object):
         self.args = args
         # We assume timestamp is a datetime.datetime object
         self.timestamp = self.args.timestamp or \
-            datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
+            datetime.datetime.now().strftime(TIMESTAMP_FMT)
         self.user = self.args.user or pwd.getpwuid(os.getuid()).pw_name
 
         self.name = self.make_run_name()
diff --git a/teuthology/suite/test/test_run_.py b/teuthology/suite/test/test_run_.py
index 1dc23e20..abe78a38 100644
--- a/teuthology/suite/test/test_run_.py
+++ b/teuthology/suite/test/test_run_.py
@@ -12,6 +12,7 @@ from io import BytesIO
 from teuthology.config import config, YamlConfig
 from teuthology.exceptions import ScheduleFailError
 from teuthology.suite import run
+from teuthology.util.time import TIMESTAMP_FMT
 
 
 class TestRun(object):
@@ -52,7 +53,7 @@ class TestRun(object):
 
     @patch('teuthology.suite.run.util.fetch_repos')
     def test_name(self, m_fetch_repos):
-        stamp = datetime.now().strftime('%Y-%m-%d_%H:%M:%S')
+        stamp = datetime.now().strftime(TIMESTAMP_FMT)
         with patch.object(run.Run, 'create_initial_config',
                           return_value=run.JobConfig()):
             name = run.Run(self.args).name
diff --git a/teuthology/util/test/test_time.py b/teuthology/util/test/test_time.py
index 33958506..c37b5f2f 100644
--- a/teuthology/util/test/test_time.py
+++ b/teuthology/util/test/test_time.py
@@ -1,11 +1,35 @@
 import pytest
 
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Type
 
 from teuthology.util import time
 
 
+@pytest.mark.parametrize(
+    ["timestamp", "result"],
+    [
+        ["1999-12-31_23:59:59", datetime(1999, 12, 31, 23, 59, 59, tzinfo=timezone.utc)],
+        ["1999-12-31_23:59", datetime(1999, 12, 31, 23, 59, 0, tzinfo=timezone.utc)],
+        ["1999-12-31T23:59:59", datetime(1999, 12, 31, 23, 59, 59, tzinfo=timezone.utc)],
+        ["1999-12-31T23:59:59+00:00", datetime(1999, 12, 31, 23, 59, 59, tzinfo=timezone.utc)],
+        ["1999-12-31T17:59:59-06:00", datetime(1999, 12, 31, 23, 59, 59, tzinfo=timezone.utc)],
+        ["2024-01-01", datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc)],
+        ["tomorrow", ValueError],
+        ["1d", ValueError],
+        ["", ValueError],
+        ["2024", ValueError],
+
+    ]
+)
+def test_parse_timestamp(timestamp: str, result: datetime | Type[Exception]):
+    if isinstance(result, datetime):
+        assert time.parse_timestamp(timestamp) == result
+    else:
+        with pytest.raises(result):
+            time.parse_timestamp(timestamp)
+
+
 @pytest.mark.parametrize(
     ["offset", "result"],
     [
diff --git a/teuthology/util/time.py b/teuthology/util/time.py
index 4fbafe84..8e0525fc 100644
--- a/teuthology/util/time.py
+++ b/teuthology/util/time.py
@@ -1,6 +1,24 @@
 import re
 
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone
+
+# When we're not using ISO format, we're using this
+TIMESTAMP_FMT = "%Y-%m-%d_%H:%M:%S"
+
+def parse_timestamp(timestamp: str) -> datetime:
+    """
+    timestamp: A string either in ISO 8601 format or TIMESTAMP_FMT.
+               If no timezone is specified, UTC is assumed.
+
+    :returns: a datetime object
+    """
+    try:
+        dt = datetime.fromisoformat(timestamp)
+    except ValueError:
+        dt = datetime.strptime(timestamp, TIMESTAMP_FMT)
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    return dt
 
 def parse_offset(offset: str) -> timedelta:
     """
-- 
2.47.3