如何测试接收可变参数的Pandas函数 测试返回None类型报错

问题原因 #

你的new_year函数采用原地修改的方式更新传入的DataFrame，没有设置显式返回值，Python中无显式return的函数默认返回None，因此你把new_year(df1)的返回值传入assert_frame_equal时，左参数就是None，自然触发类型不匹配的报错。

解决方法 #

有两种可选方案，你可以根据业务需求选择：

方案1：不修改业务代码，调整测试逻辑 #

既然函数是原地修改入参，那调用完函数后直接对比修改后的入参即可，调整测试代码如下：

import pandas as pd
from pandas._testing import assert_frame_equal
from prog import new_year

def test_new_year():
    df1 = pd.DataFrame({'Name':['Jack', 'Jill'], 'Class':[4, 5]})
    df2 = pd.DataFrame({'Name':['Jack', 'Jill'], 'Class':[5, 6]})
    # 先调用函数修改df1
    new_year(df1)
    # 再对比修改后的df1和预期结果df2
    assert_frame_equal(df1, df2)

方案2：修改业务代码，让函数返回修改后的DataFrame #

如果你希望函数有返回值方便链式调用，可以给业务函数加返回语句，修改prog.py中的new_year函数：

def new_year(df):
    df.Class += 1
    # 新增返回语句
    return df

修改后你原来写的测试代码不需要调整就能正常运行。

内容的提问来源于stack exchange，提问作者flywire