use a better test

rapidsai · Jun 6, 2024 · 0905a28 · 0905a28
1 parent 5be89d5
commit 0905a28
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/python/dask_cudf/dask_cudf/io/tests/test_parquet.py b/python/dask_cudf/dask_cudf/io/tests/test_parquet.py
@@ -610,10 +610,9 @@ def test_timezone_column(tmpdir):
         }
     )
     pdf.to_parquet(path)
-    # cudf.read_parquet does not support reading timezone aware types yet, so check dtypes
-    got = dask_cudf.read_parquet(path).dtypes
-    expected = pd.Series(
-        {"time": pd.DatetimeTZDtype("ns", "UTC"), "x": np.dtype(np.int64)},
-        dtype=object,
-    )
-    pd.testing.assert_series_equal(got, expected)
+    got = dask_cudf.read_parquet(path)
+    # cudf.read_parquet does not support reading timezone aware types yet
+    assert got["time"].dtype == pd.DatetimeTZDtype("ns", "UTC")
+    got["time"] = got["time"].astype("datetime64[ns]")
+    expected = cudf.read_parquet(path)
+    dd.assert_eq(got, expected)