updated testing logic

Author: arthurlw

pandas/tests/extension/test_arrow.py

@@ -64,6 +64,7 @@
     is_string_dtype,
     is_unsigned_integer_dtype,
 )
+from pandas.core.dtypes.common import is_timedelta64_dtype
 from pandas.tests.extension import base
 
 pa = pytest.importorskip("pyarrow")
@@ -281,15 +282,33 @@ def test_compare_scalar(self, data, comparison_op):
     @pytest.mark.parametrize("na_action", [None, "ignore"])
     def test_map(self, data_missing, na_action):
         if data_missing.dtype.kind in "mM":
-            result = pd.Series(
-                np.asarray(
-                    data_missing.map(lambda x: x, na_action=na_action), dtype="int64"
-                )
-            )
-            expected = pd.Series(
-                data_missing.to_numpy().astype(result.dtype).view("int64")
-            )
-            tm.assert_series_equal(result, expected, check_dtype=False)
+            mapped = data_missing.map(lambda x: x, na_action=na_action)
+            result = pd.Series(mapped)
+            expected = pd.Series(data_missing.to_numpy())
+
+            orig_dtype = expected.dtype
+
+            if result.dtype == "float64" and (
+                is_datetime64_any_dtype(orig_dtype)
+                or is_timedelta64_dtype(orig_dtype)
+                or isinstance(orig_dtype, pd.DatetimeTZDtype)
+            ):
+                result = result.astype(orig_dtype)
+
+            if isinstance(orig_dtype, pd.DatetimeTZDtype):
+                pass
+            elif is_datetime64_any_dtype(orig_dtype):
+                result = result.astype("datetime64[ns]").astype("int64")
+                expected = expected.astype("datetime64[ns]").astype("int64")
+                result = pd.Series(result)
+                expected = pd.Series(expected)
+            elif is_timedelta64_dtype(orig_dtype):
+                result = result.astype("timedelta64[ns]")
+                expected = expected.astype("timedelta64[ns]")
+
+
+            tm.assert_series_equal(result, expected, check_dtype=False, check_exact=False)
+
         else:
             result = data_missing.map(lambda x: x, na_action=na_action)
             if data_missing.dtype == "float32[pyarrow]":