Skip to content

Commit d859236

Browse files
CLN: Refactor tests to adress PR feedback
1 parent 1beecf1 commit d859236

File tree

2 files changed

+11
-15
lines changed

2 files changed

+11
-15
lines changed

pandas/tests/reshape/test_pivot.py

+4-8
Original file line numberDiff line numberDiff line change
@@ -2858,8 +2858,9 @@ def test_pivot_with_pyarrow_categorical(self):
28582858
# GH#53051
28592859
pa = pytest.importorskip("pyarrow")
28602860

2861+
data = {"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]}
28612862
df = (
2862-
DataFrame({"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]})
2863+
DataFrame(data)
28632864
.astype({"string_column": "category", "number_column": "float32"})
28642865
.astype(
28652866
{
@@ -2871,13 +2872,8 @@ def test_pivot_with_pyarrow_categorical(self):
28712872

28722873
df = df.pivot(columns=["string_column"], values=["number_column"])
28732874

2874-
multi_index = MultiIndex.from_arrays(
2875-
[["number_column", "number_column", "number_column"], ["A", "B", "C"]],
2876-
names=(None, "string_column"),
2877-
)
2878-
df_expected = DataFrame(
2879-
[[1.0, np.nan, np.nan], [np.nan, 2.0, np.nan], [np.nan, np.nan, 3.0]],
2880-
columns=multi_index,
2875+
df_expected = DataFrame(data).pivot(
2876+
columns=["string_column"], values=["number_column"]
28812877
)
28822878
tm.assert_frame_equal(
28832879
df, df_expected, check_dtype=False, check_column_type=False

pandas/tests/test_multilevel.py

+7-7
Original file line numberDiff line numberDiff line change
@@ -324,8 +324,9 @@ def test_multiindex_with_pyarrow_categorical(self):
324324
# GH#53051
325325
pa = pytest.importorskip("pyarrow")
326326

327+
data = {"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]}
327328
df = (
328-
DataFrame({"string_column": ["A", "B", "C"], "number_column": [1, 2, 3]})
329+
DataFrame(data)
329330
.astype({"string_column": "category", "number_column": "float32"})
330331
.astype(
331332
{
@@ -337,13 +338,12 @@ def test_multiindex_with_pyarrow_categorical(self):
337338

338339
df = df.set_index(["string_column", "number_column"])
339340

340-
df_expected = DataFrame(
341-
index=MultiIndex.from_arrays(
342-
[["A", "B", "C"], [1, 2, 3]], names=["string_column", "number_column"]
343-
)
344-
)
341+
df_expected = DataFrame(data).set_index(["string_column", "number_column"])
345342
tm.assert_frame_equal(
346-
df, df_expected, check_dtype=False, check_column_type=False
343+
df,
344+
df_expected,
345+
check_dtype=False,
346+
check_index_type=False,
347347
)
348348

349349

0 commit comments

Comments
 (0)