Fix test

dlt-hub · Sep 23, 2023 · 83e66f7 · 83e66f7
1 parent 4021987
commit 83e66f7
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/dlt/common/time.py b/dlt/common/time.py
@@ -1,5 +1,5 @@
 import contextlib
-from typing import Any, Optional, Union, overload  # noqa
+from typing import Any, Optional, Union, overload, TypeVar  # noqa
 import datetime  # noqa: I251
 
 from dlt.common.pendulum import pendulum, timedelta
@@ -148,5 +148,7 @@ def to_seconds(td: Optional[TimedeltaSeconds]) -> Optional[float]:
     return td
 
 
-def reduce_pendulum_datetime_precision(value: pendulum.DateTime, microsecond_precision: int) -> pendulum.DateTime:
-    return value.set(microsecond=value.microsecond // 10**(6 - microsecond_precision) * 10**(6 - microsecond_precision))  # type: ignore
+T = TypeVar("T", bound=Union[pendulum.DateTime, pendulum.Time])
+
+def reduce_pendulum_datetime_precision(value: T, microsecond_precision: int) -> T:
+    return value.replace(microsecond=value.microsecond // 10**(6 - microsecond_precision) * 10**(6 - microsecond_precision))  # type: ignore
diff --git a/tests/cases.py b/tests/cases.py
@@ -7,7 +7,7 @@
 from dlt.common.data_types import TDataType
 from dlt.common.typing import StrAny
 from dlt.common.wei import Wei
-from dlt.common.time import ensure_pendulum_datetime, reduce_pendulum_datetime_precision
+from dlt.common.time import ensure_pendulum_datetime, reduce_pendulum_datetime_precision, ensure_pendulum_time
 from dlt.common.schema import TColumnSchema, TTableSchemaColumns
 
 
@@ -274,6 +274,21 @@ def assert_all_data_types_row(
             ensure_pendulum_datetime(expected_rows["col4"]),  # type: ignore[arg-type]
             timestamp_precision
         )
+    if "col4_precision" in expected_rows:
+        parsed_date = pendulum.instance(db_mapping["col4_precision"])
+        db_mapping["col4_precision"] = reduce_pendulum_datetime_precision(parsed_date, 3)
+        expected_rows['col4_precision'] = reduce_pendulum_datetime_precision(
+            ensure_pendulum_datetime(expected_rows["col4_precision"]),  # type: ignore[arg-type]
+            3
+        )
+
+    if "col11_precision" in expected_rows:
+        parsed_time = ensure_pendulum_time(db_mapping["col11_precision"])
+        db_mapping["col11_precision"] = reduce_pendulum_datetime_precision(parsed_time, 3)
+        expected_rows['col11_precision'] = reduce_pendulum_datetime_precision(
+            ensure_pendulum_time(expected_rows["col11_precision"]),  # type: ignore[arg-type]
+            3
+        )
 
     # binary column
     if "col7" in db_mapping: