From 1b2fadb953dd7192d33b3298e26573c0a7cb40f0 Mon Sep 17 00:00:00 2001
From: anders-albert <anders0albert@gmail.com>
Date: Tue, 3 Sep 2024 15:33:09 +0200
Subject: [PATCH] fix: handle missing value

---
 .../data_modeling/typed_instances.py          |  5 ++--
 cognite/client/utils/_time.py                 | 21 +++++--------
 tests/tests_unit/test_utils/test_time.py      | 30 +++++++++++++++++++
 3 files changed, 40 insertions(+), 16 deletions(-)

diff --git a/cognite/client/data_classes/data_modeling/typed_instances.py b/cognite/client/data_classes/data_modeling/typed_instances.py
index e6a064e765..dbf7e973b7 100644
--- a/cognite/client/data_classes/data_modeling/typed_instances.py
+++ b/cognite/client/data_classes/data_modeling/typed_instances.py
@@ -3,6 +3,7 @@
 import inspect
 from abc import ABC
 from collections.abc import Iterable
+from datetime import date
 from typing import TYPE_CHECKING, Any, cast
 
 from typing_extensions import Self
@@ -19,7 +20,7 @@
     _serialize_property_value,
 )
 from cognite.client.utils._text import to_camel_case
-from cognite.client.utils._time import date_str_to_date, timestamp_str_to_datetime
+from cognite.client.utils._time import timestamp_str_to_datetime
 
 if TYPE_CHECKING:
     from cognite.client import CogniteClient
@@ -306,7 +307,7 @@ def _deserialize_value(value: Any, parameter: inspect.Parameter) -> Any:
     if "datetime" in annotation and isinstance(value, str):
         return timestamp_str_to_datetime(value)
     elif "date" in annotation and isinstance(value, str):
-        return date_str_to_date(value)
+        return date.fromisoformat(value)
     elif DirectRelationReference.__name__ in annotation and isinstance(value, dict):
         return DirectRelationReference.load(value)
     elif NodeId.__name__ in annotation and isinstance(value, dict):
diff --git a/cognite/client/utils/_time.py b/cognite/client/utils/_time.py
index 776d145ceb..370ba68996 100644
--- a/cognite/client/utils/_time.py
+++ b/cognite/client/utils/_time.py
@@ -9,7 +9,7 @@
 import time
 from abc import ABC, abstractmethod
 from contextlib import suppress
-from datetime import date, datetime, timedelta, timezone
+from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, cast, overload
 
 from cognite.client.utils._importing import local_import
@@ -187,19 +187,12 @@ def timestamp_str_to_datetime(timestamp: str) -> datetime:
     Returns:
         datetime: Datetime object
     """
-    return datetime.fromisoformat(timestamp)
-
-
-def date_str_to_date(date_str: str) -> date:
-    """Converts a date string in ISO 8601 format to a datetime object.
-
-    Args:
-        date_str (str): Date string in ISO 8601 format
-
-    Returns:
-        date: Datetime object
-    """
-    return datetime.fromisoformat(date_str)
+    try:
+        return datetime.fromisoformat(timestamp)
+    except ValueError:
+        # Typically hits if the timestamp is missing milliseconds
+        # For example, "2021-01-01T00:00:00.17+00:00", i.e. missing the last digit
+        return datetime.strptime(timestamp, "%Y-%m-%dT%H:%M:%S.%f%z")
 
 
 def split_granularity_into_quantity_and_normalized_unit(granularity: str) -> tuple[int, str]:
diff --git a/tests/tests_unit/test_utils/test_time.py b/tests/tests_unit/test_utils/test_time.py
index b55250d308..63e07519f2 100644
--- a/tests/tests_unit/test_utils/test_time.py
+++ b/tests/tests_unit/test_utils/test_time.py
@@ -27,6 +27,7 @@
     parse_str_timezone,
     parse_str_timezone_offset,
     split_time_range,
+    timestamp_str_to_datetime,
     timestamp_to_ms,
     to_fixed_utc_intervals,
     to_pandas_freq,
@@ -230,6 +231,35 @@ def test_negative(self, t):
             timestamp_to_ms(t)
 
 
+class TestTimestampStrToDatetime:
+    @pytest.mark.parametrize(
+        "timestamp_str, expected",
+        [
+            ("2021-01-01T00:00:00.000+00:00", datetime(2021, 1, 1, 0, 0, 0, 0, tzinfo=timezone.utc)),
+            ("2021-01-01T00:00:00.000+01:00", datetime(2021, 1, 1, 0, 0, 0, 0, tzinfo=timezone(timedelta(hours=1)))),
+            (
+                "2021-01-01T00:00:00.000+01:15",
+                datetime(2021, 1, 1, 0, 0, 0, 0, tzinfo=timezone(timedelta(hours=1, minutes=15))),
+            ),
+            (
+                "2021-01-01T00:00:00.000-01:15",
+                datetime(2021, 1, 1, 0, 0, 0, 0, tzinfo=timezone(timedelta(hours=-1, minutes=-15))),
+            ),
+            ("2024-09-03T09:36:01.17+00:00", datetime(2024, 9, 3, 9, 36, 1, 170000, tzinfo=timezone.utc)),
+        ],
+    )
+    def test_valid_timestamp_str(self, timestamp_str: str, expected: datetime):
+        assert expected == timestamp_str_to_datetime(timestamp_str)
+
+    @pytest.mark.parametrize(
+        "timestamp_str",
+        ["2021-01-01T00:00:00.000", "2021-01-01T00:00:00.000+01:15:12", "2021-01-01T00:00:00.000+01:15:12:13"],
+    )
+    def test_invalid_timestamp_str(self, timestamp_str):
+        with pytest.raises(TypeError, match="Invalid timestamp format"):
+            timestamp_str_to_datetime(timestamp_str)
+
+
 class TestGranularityToMs:
     @pytest.mark.parametrize(
         "granularity, expected_ms",