unit test on _set_tz

Author: chaoyihu

pandas/tests/io/pytables/test_read.py

@@ -317,14 +317,3 @@ def test_read_infer_string(tmp_path, setup_path):
         columns=Index(["a"], dtype="string[pyarrow_numpy]"),
     )
     tm.assert_frame_equal(result, expected)
-
-
-def test_hdfstore_read_datetime64_unit_s(tmp_path, setup_path):
-    # Fix issue 59004: HDFStore doesn't save datetime64[s] right
-    df_s = DataFrame(["2001-01-01", "2002-02-02"], dtype="datetime64[s]")
-    path = tmp_path / setup_path
-    with HDFStore(path, mode="w") as store:
-        store.put("df_s", df_s)
-    with HDFStore(path, mode="r") as store:
-        df_fromstore = store.get("df_s")
-    tm.assert_frame_equal(df_s, df_fromstore)

pandas/tests/io/pytables/test_timezones.py

@@ -23,6 +23,8 @@
     ensure_clean_store,
 )
 
+from pandas.io.pytables import _set_tz
+
 
 def _compare_with_tz(a, b):
     tm.assert_frame_equal(a, b)
@@ -345,3 +347,22 @@ def test_read_with_where_tz_aware_index(tmp_path, setup_path):
         store.append(key, expected, format="table", append=True)
     result = pd.read_hdf(path, key, where="DATE > 20151130")
     tm.assert_frame_equal(result, expected)
+
+
+def test_set_tz_datetime64_unit(tmp_path, setup_path):
+    # GH 59004
+    # https://numpy.org/doc/stable/reference/generated/numpy.datetime_data.html
+    tz = None
+    example_dtype = ["datetime64[s]", "datetime64[25s]", "datetime64[ns]"]
+    example_values = [
+        np.array([978307200, 1012608000], dtype=np.int64),
+        np.array([978307200, 1012608000], dtype=np.int64),
+        np.array([978307200000000000, 1012608000000000000], dtype=np.int64),
+    ]
+    for dtype, values in zip(example_dtype, example_values):
+        tm.assert_datetime_array_equal(
+            _set_tz(values, tz, dtype),
+            pd.arrays.DatetimeArray._from_sequence(
+                DatetimeIndex(["2001-01-01", "2002-02-02"], dtype=dtype)
+            ),
+        )