Fix test case and move import

Author: ABizzinotto

pandas/tests/strings/test_extract.py

@@ -2,17 +2,17 @@
 import re
 
 import numpy as np
-import pyarrow as pa
 import pytest
 
+from pandas.core.dtypes.dtypes import ArrowDtype
+
 from pandas import (
     DataFrame,
     Index,
     MultiIndex,
     Series,
     _testing as tm,
 )
-from pandas.core.dtypes.dtypes import ArrowDtype
 
 
 def test_extract_expand_kwarg_wrong_type_raises(any_string_dtype):
@@ -710,16 +710,10 @@ def test_extractall_same_as_extract_subject_index(any_string_dtype):
     tm.assert_frame_equal(extract_one_noname, no_match_index)
 
 
-@pytest.mark.parametrize(
-    "data, expected_dtype",
-    [
-        (Series(["abc", "ab"], dtype=ArrowDtype(pa.string())), "string[pyarrow]"),
-        (Series(["abc", "ab"], dtype="string"), "string[python]"),
-        (Series(["abc", "ab"]), "object"),
-    ]
-)
-def test_extractall_preserves_dtype(data, expected_dtype):
+def test_extractall_preserves_dtype():
     # Ensure that when extractall is called on a series with specific dtypes set, that
     # the dtype is preserved in the resulting DataFrame's column.
-    result = data.str.extractall("(ab)")
-    assert result.dtypes[0] == expected_dtype
+    import pyarrow as pa
+
+    result = Series(["abc", "ab"], dtype=ArrowDtype(pa.string())).str.extractall("(ab)")
+    assert result.dtypes[0] == "string[pyarrow]"