From a10b643c2d8565e7017645e9e0abeec79f5f25f1 Mon Sep 17 00:00:00 2001
From: Jeff Mattson <jmattson@sei.cmu.edu>
Date: Thu, 13 Jun 2024 14:04:13 -0400
Subject: [PATCH] support timestamps without microseconds

---
 takproto/constants.py   |  4 ++--
 takproto/functions.py   |  6 +++++-
 tests/test_functions.py | 13 +++++++++++++
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/takproto/constants.py b/takproto/constants.py
index 279bde6..3cbf3e6 100644
--- a/takproto/constants.py
+++ b/takproto/constants.py
@@ -34,8 +34,8 @@
 DEFAULT_PROTO_HEADER = bytearray(b"\xbf")
 DEFAULT_MESH_HEADER = bytearray(b"\xbf\x01\xbf")
 
-W3C_XML_DATETIME: str = "%Y-%m-%dT%H:%M:%S.%fZ"
-ISO_8601_UTC = W3C_XML_DATETIME  # Issue 7: Not technically correct.
+W3C_XML_DATETIME: str = "%Y-%m-%dT%H:%M:%SZ"
+ISO_8601_UTC: str = "%Y-%m-%dT%H:%M:%S.%fZ"
 
 
 class TAKProtoVer(Enum):
diff --git a/takproto/functions.py b/takproto/functions.py
index a18c0cb..cca4b51 100644
--- a/takproto/functions.py
+++ b/takproto/functions.py
@@ -35,6 +35,7 @@
 
 from takproto.constants import (
     ISO_8601_UTC,
+    W3C_XML_DATETIME,
     DEFAULT_MESH_HEADER,
     DEFAULT_PROTO_HEADER,
     TAKProtoVer,
@@ -70,7 +71,10 @@ def parse_stream(msg):
 
 def format_time(time: str) -> int:
     """Format timestamp as microseconds."""
-    s_time = datetime.strptime(time + "+0000", ISO_8601_UTC + "%z")
+    try:
+        s_time = datetime.strptime(time + "+0000", W3C_XML_DATETIME + "%z")
+    except ValueError:
+        s_time = datetime.strptime(time + "+0000", ISO_8601_UTC + "%z")
     return int(s_time.timestamp() * 1000)
 
 
diff --git a/tests/test_functions.py b/tests/test_functions.py
index a20c85c..2ffd48c 100644
--- a/tests/test_functions.py
+++ b/tests/test_functions.py
@@ -47,6 +47,19 @@ def test_format_timestamp(self):
         )
         self.assertEqual(time2, t_time)
 
+    def test_format_timestamp_without_subseconds(self):
+        """Test formatting timestamp to and from Protobuf format."""
+        t_time = "2020-02-08T18:10:44Z"
+        t_ts = 1581185444000
+        ts = takproto.format_time(t_time)
+        self.assertEqual(ts, t_ts)
+
+        t_ts2 = t_ts / 1000
+        time2 = datetime.fromtimestamp(t_ts2, timezone.utc).strftime(
+            "%Y-%m-%dT%H:%M:%SZ"
+        )
+        self.assertEqual(time2, t_time)
+
     def test_xml2proto_default(self):
         """Test encoding XML string as Protobuf bytearray."""
         t_xml = """<?xml version='1.0' encoding='UTF-8' standalone='yes'?>