Adjust tests in internals folder for arrow string option (#56140)

pandas-dev · Nov 26, 2023 · 22d3c9c · 22d3c9c
1 parent 472516c
commit 22d3c9c
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/pandas/tests/internals/test_internals.py b/pandas/tests/internals/test_internals.py
@@ -591,7 +591,7 @@ def test_astype(self, t):
         else:
             assert tmgr.iget(3).dtype.type == t
 
-    def test_convert(self):
+    def test_convert(self, using_infer_string):
         def _compare(old_mgr, new_mgr):
             """compare the blocks, numeric compare ==, object don't"""
             old_blocks = set(old_mgr.blocks)
@@ -626,9 +626,10 @@ def _compare(old_mgr, new_mgr):
         mgr.iset(1, np.array(["2."] * N, dtype=np.object_))
         mgr.iset(2, np.array(["foo."] * N, dtype=np.object_))
         new_mgr = mgr.convert(copy=True)
-        assert new_mgr.iget(0).dtype == np.object_
-        assert new_mgr.iget(1).dtype == np.object_
-        assert new_mgr.iget(2).dtype == np.object_
+        dtype = "string[pyarrow_numpy]" if using_infer_string else np.object_
+        assert new_mgr.iget(0).dtype == dtype
+        assert new_mgr.iget(1).dtype == dtype
+        assert new_mgr.iget(2).dtype == dtype
         assert new_mgr.iget(3).dtype == np.int64
         assert new_mgr.iget(4).dtype == np.float64
 
@@ -639,9 +640,9 @@ def _compare(old_mgr, new_mgr):
         mgr.iset(1, np.array(["2."] * N, dtype=np.object_))
         mgr.iset(2, np.array(["foo."] * N, dtype=np.object_))
         new_mgr = mgr.convert(copy=True)
-        assert new_mgr.iget(0).dtype == np.object_
-        assert new_mgr.iget(1).dtype == np.object_
-        assert new_mgr.iget(2).dtype == np.object_
+        assert new_mgr.iget(0).dtype == dtype
+        assert new_mgr.iget(1).dtype == dtype
+        assert new_mgr.iget(2).dtype == dtype
         assert new_mgr.iget(3).dtype == np.int32
         assert new_mgr.iget(4).dtype == np.bool_
         assert new_mgr.iget(5).dtype.type, np.datetime64