如何对数据框进行单元测试或使用断言？

Question

我想对我的代码进行单元测试，以确保代码的输出是正确的。

我也从 StackOverflow 尝试过，但没有工作：

input = pd.DataFrame.from_dict({
    'rev': [0],
    'price': [0]
})
expected = {
    'data': ['cool']
}

assert_dict_equal(expected, data(input).to_dict(),
                  "oops, there's a bug...")


def temperature(row):
    rev = row['rev']
    price = row['price']
    group = row['group']
    if group == 'error':
        return 'error'
    elif revenue > 2 * price:
        return 'cold'
    elif revenue >= price:
        return 'cool'
    elif revenue < 0.5 * price:
        return 'hot'
    elif revenue < price:
        return 'lukewarm'
    else:
        float('NA')

data['temp'] = sample.apply(temperature, axis=1)

assert temperature({"group": 'T2_Y2', "rev": 0, "price": 0}, 
                        "group", "rev", "price") == 'cool'

assert temperature({"group": 'T2_Y2', "rev": 30, "price": 10}, 
                        "group", "rev", "price") == 'cold'

assert temperature({"group": 'T2_Y2', "rev": 3, "price": 10}, 
                        "group", "rev", "price") == 'hot'

当输出给出错误时，预期的结果应该返回而没有错误。

TypeError: temperature() 接受 1 个位置参数，但给出了 4 个

Answer 1

错误很明确：您试图在断言中使用 4 个位置参数调用 temperature 函数，而它只需要 1 个。用这些替换您的断言：

assert temperature({"group": 'T2_Y2', "rev": 0, "price": 0}) == 'cool'
assert temperature({"group": 'T2_Y2', "rev": 30, "price": 10}) == 'cold'
assert temperature({"group": 'T2_Y2', "rev": 3, "price": 10}) == 'hot'

编辑：根据您的评论，将您的功能更改为以下功能：

def temperature(row, group_col, rev_col, price_col):
    rev = row[rev_col]
    price = row[price_col]
    group = row[group_col]
    if group == 'error':
        return 'error'
    elif revenue > 2 * price:
        return 'cold'
    elif revenue >= price:
        return 'cool'
    elif revenue < 0.5 * price:
        return 'hot'
    elif revenue < price:
        return 'lukewarm'
    else:
        float('NA')

然后保持你的断言不变，它应该可以工作