non portable type resolution solution

kuzudb · Mar 15, 2024 · f2fae8f · f2fae8f
1 parent 3828ace
commit f2fae8f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tools/python_api/test/test_df_pyarrow.py b/tools/python_api/test/test_df_pyarrow.py
@@ -133,7 +133,9 @@ def test_pyarrow_time(conn_db_readonly : ConnDB) -> None:
             assert expected.to_pytimedelta() == actual.to_pytimedelta()
     for colname in ['col4', 'col5', 'col6', 'col7', 'col8']:
         for expected, actual in zip(df[colname], result[colname]):
-            assert expected.to_pydatetime() == actual.to_pydatetime()
+            tmp1 = expected if type(expected) is datetime else expected.to_pydatetime()
+            tmp2 = actual if type(actual) is datetime else actual.to_pydatetime()
+            assert tmp1 == tmp2
     for colname in ['col9', 'col10']:
         for expected, actual in zip(df[colname], result[colname]):
             assert datetime.combine(expected, datetime.min.time()) == actual.to_pydatetime()