Skip to content

Commit

Permalink
used pytest.mark.parametrize for expected_dtype
Browse files Browse the repository at this point in the history
  • Loading branch information
AbhishekChaudharii committed Nov 21, 2024
1 parent 1f25150 commit 984b7d6
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions pandas/tests/reshape/test_pivot.py
Original file line number Diff line number Diff line change
Expand Up @@ -2376,15 +2376,16 @@ def test_pivot_table_with_margins_and_numeric_columns(self):

tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize("dtype", ["Int64", "int64"])
def test_pivot_ea_dtype_dropna(self, dropna, dtype):
@pytest.mark.parametrize(
"dtype,expected_dtype", [("Int64", "Float64"), ("int64", "float64")]
)
def test_pivot_ea_dtype_dropna(self, dropna, dtype, expected_dtype):
# GH#47477
# GH#47971
df = DataFrame({"x": "a", "y": "b", "age": Series([20, 40], dtype=dtype)})
result = df.pivot_table(
index="x", columns="y", values="age", aggfunc="mean", dropna=dropna
)
expected_dtype = "float64" if dtype == "int64" else "Float64"
expected = DataFrame(
[[30]],
index=Index(["a"], name="x"),
Expand Down

0 comments on commit 984b7d6

Please sign in to comment.