pandas panda单元测试Assert错误: Dataframe 索引不同

d7v8vwbk  于 2022-12-09  发布在  其他
关注(0)|答案(1)|浏览(190)

我想测试这个函数:

def filter_df(df, column_name: str, skill: List):
    return df.query(f"{column_name} in {skill}")

这是我的考验:

def test_filter_df():
    df = pd.DataFrame({"col1": ["sap", "hi", "abc"], "col2": [3, 4, 4]})
    expected = pd.DataFrame({"col1": ["hi", "abc"], "col2": [4, 4]})
    assert_frame_equal(filter_df(df, "col1", ["hi", "abc"]), expected)

我得到一个assert_frame_equal(filter_df(df, "col1", ["hi", "abc"]), expected)错误,但我不明白为什么 Dataframe 不相同。

w8f9ii69

w8f9ii691#

您需要重置filter_df中的索引:

df.query(f"{column_name} in {skill}").reset_index(drop=True)

此时,返回的DF具有给定行的原始索引,在您的情况下为1,2,而不是expected DF中的0,1
或者,如果这是函数的预期行为,请编辑expected DF以获得正确的索引

相关问题